Update app.py

app.py

@@ -25,7 +25,7 @@ from datasets import load_dataset
 import tempfile
 
 st.title("SQL-RAG Using CrewAI 🚀")
-st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
+st.write("Analyze datasets using natural language queries.")
 
 # Initialize LLM
 llm = None
@@ -86,106 +86,172 @@ 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.
-    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"
-    }}
-    """
-    response = llm.generate(prompt)
-    try:
-        return json.loads(response)
-    except json.JSONDecodeError:
-        st.error("⚠️ GPT-4o failed to generate a valid suggestion.")
-        return None
 
-# Dynamically generate Plotly visualizations based on GPT-4o suggestions
-def generate_visualization(suggestion, df):
-    chart_type = suggestion.get("chart_type", "bar").lower()
-    x_axis = suggestion.get("x_axis")
-    y_axis = suggestion.get("y_axis")
-    group_by = suggestion.get("group_by")
+# Helper Function for Validation
+def is_valid_suggestion(suggestion):
+    chart_type = suggestion.get("chart_type", "").lower()
 
-    # Ensure required inputs are available
-    if not x_axis or not y_axis:
-        st.warning("⚠️ GPT-4o did not provide enough information for the visualization.")
-        return None
+    if chart_type in ["bar", "line", "box", "scatter"]:
+        return all(k in suggestion for k in ["chart_type", "x_axis", "y_axis"])
 
-    # Dynamically select the Plotly function
-    plotly_function = getattr(px, chart_type, None)
-
-    # Handle unsupported chart types gracefully
-    if not plotly_function:
-        st.warning(f"⚠️ Unsupported chart type '{chart_type}' suggested by GPT-4o.")
-        return None
-
-    # Prepare dynamic parameters for Plotly function
-    plot_args = {
-        "data_frame": df,
-        "x": x_axis,
-        "y": y_axis,
-    }
-    if group_by:
-        plot_args["color"] = group_by
+    elif chart_type == "pie":
+        return all(k in suggestion for k in ["chart_type", "x_axis"])
 
-    try:
-        # Generate the dynamic visualization
-        fig = plotly_function(**plot_args)
-        fig.update_layout(
-            title=f"{chart_type.title()} Plot of {y_axis.replace('_', ' ').title()} by {x_axis.replace('_', ' ').title()}",
-            xaxis_title=x_axis.replace('_', ' ').title(),
-            yaxis_title=y_axis.replace('_', ' ').title(),
-        )
-        return fig
+    elif chart_type == "heatmap":
+        return all(k in suggestion for k in ["chart_type", "x_axis", "y_axis"])
 
-    except Exception as e:
-        st.error(f"⚠️ Failed to generate visualization: {e}")
-        return None
-
-# Function to create TXT file
-def create_text_report_with_viz_temp(report, conclusion, visualizations):
-    content = f"### Analysis Report\n\n{report}\n\n### Visualizations\n"
+    else:
+        return False
+        
+def ask_gpt4o_for_visualization(query, df, llm, retries=2):
+    import json
 
-    for i, fig in enumerate(visualizations, start=1):
-        fig_title = fig.layout.title.text if fig.layout.title.text else f"Visualization {i}"
-        x_axis = fig.layout.xaxis.title.text if fig.layout.xaxis.title.text else "X-axis"
-        y_axis = fig.layout.yaxis.title.text if fig.layout.yaxis.title.text else "Y-axis"
+    # Identify numeric and categorical columns
+    numeric_columns = df.select_dtypes(include='number').columns.tolist()
+    categorical_columns = df.select_dtypes(exclude='number').columns.tolist()
 
-        content += f"\n{i}. {fig_title}\n"
-        content += f"   - X-axis: {x_axis}\n"
-        content += f"   - Y-axis: {y_axis}\n"
+    # Prompt with Dataset-Specific, Query-Based Examples
+    prompt = f"""
+    Analyze the following query and suggest the most suitable visualization(s) using the dataset.
+    **Query:** "{query}"
+    **Dataset Overview:**  
+    - **Numeric Columns (for Y-axis):** {', '.join(numeric_columns) if numeric_columns else 'None'}  
+    - **Categorical Columns (for X-axis or grouping):** {', '.join(categorical_columns) if categorical_columns else 'None'}  
+    Suggest visualizations in this exact JSON format:
+    [
+      {{
+        "chdart_type": "bar/box/line/scatter/pie/heatmap",
+        "x_axis": "categorical_or_time_column",
+        "y_axis": "numeric_column",
+        "group_by": "optional_column_for_grouping",
+        "title": "Title of the chart",
+        "description": "Why this chart is suitable"
+      }}
+    ]
+    **Query-Based Examples:**  
+    - **Query:** "What is the salary distribution across different job titles?"  
+      **Suggested Visualization:**  
+      {{
+        "chart_type": "box",
+        "x_axis": "job_title",
+        "y_axis": "salary_in_usd",
+        "group_by": "experience_level",
+        "title": "Salary Distribution by Job Title and Experience",
+        "description": "A box plot to show how salaries vary across different job titles and experience levels."
+      }}
+    - **Query:** "Show the average salary by company size and employment type."  
+      **Suggested Visualizations:**  
+      [
+        {{
+          "chart_type": "bar",
+          "x_axis": "company_size",
+          "y_axis": "salary_in_usd",
+          "group_by": "employment_type",
+          "title": "Average Salary by Company Size and Employment Type",
+          "description": "A grouped bar chart comparing average salaries across company sizes and employment types."
+        }},
+        {{
+          "chart_type": "heatmap",
+          "x_axis": "company_size",
+          "y_axis": "salary_in_usd",
+          "group_by": "employment_type",
+          "title": "Salary Heatmap by Company Size and Employment Type",
+          "description": "A heatmap showing salary concentration across company sizes and employment types."
+        }}
+      ]
+    - **Query:** "How has the average salary changed over the years?"  
+      **Suggested Visualization:**  
+      {{
+        "chart_type": "line",
+        "x_axis": "work_year",
+        "y_axis": "salary_in_usd",
+        "group_by": "experience_level",
+        "title": "Average Salary Trend Over Years",
+        "description": "A line chart showing how the average salary has changed across different experience levels over the years."
+      }}
+    - **Query:** "What is the employee distribution by company location?"  
+      **Suggested Visualization:**  
+      {{
+        "chart_type": "pie",
+        "x_axis": "company_location",
+        "y_axis": null,
+        "group_by": null,
+        "title": "Employee Distribution by Company Location",
+        "description": "A pie chart showing the distribution of employees across company locations."
+      }}
+    - **Query:** "Is there a relationship between remote work ratio and salary?"  
+      **Suggested Visualization:**  
+      {{
+        "chart_type": "scatter",
+        "x_axis": "remote_ratio",
+        "y_axis": "salary_in_usd",
+        "group_by": "experience_level",
+        "title": "Remote Work Ratio vs Salary",
+        "description": "A scatter plot to analyze the relationship between remote work ratio and salary."
+      }}
+    - **Query:** "Which job titles have the highest salaries across regions?"  
+      **Suggested Visualization:**  
+      {{
+        "chart_type": "heatmap",
+        "x_axis": "job_title",
+        "y_axis": "employee_residence",
+        "group_by": null,
+        "title": "Salary Heatmap by Job Title and Region",
+        "description": "A heatmap showing the concentration of high-paying job titles across regions."
+      }}
+    Only suggest visualizations that logically match the query and dataset.
+    """
 
-        if fig.data:
-            trace_types = set(trace.type for trace in fig.data)
-            content += f"   - Chart Type(s): {', '.join(trace_types)}\n"
-        else:
-            content += "   - No data available in this visualization.\n"
+    for attempt in range(retries + 1):
+        try:
+            response = llm.generate(prompt)
+            suggestions = json.loads(response)
+
+            if isinstance(suggestions, list):
+                valid_suggestions = [s for s in suggestions if is_valid_suggestion(s)]
+                if valid_suggestions:
+                    return valid_suggestions
+                else:
+                    st.warning("⚠️ GPT-4o did not suggest valid visualizations.")
+                    return None
+
+            elif isinstance(suggestions, dict):
+                if is_valid_suggestion(suggestions):
+                    return [suggestions]
+                else:
+                    st.warning("⚠️ GPT-4o's suggestion is incomplete or invalid.")
+                    return None
+
+        except json.JSONDecodeError:
+            st.warning(f"⚠️ Attempt {attempt + 1}: GPT-4o returned invalid JSON.")
+        except Exception as e:
+            st.error(f"⚠️ Error during GPT-4o call: {e}")
 
-    content += f"\n\n\n{conclusion}"
+        if attempt < retries:
+            st.info("🔄 Retrying visualization suggestion...")
 
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode='w', encoding='utf-8') as temp_txt:
-        temp_txt.write(content)
-        return temp_txt.name
+    st.error("❌ Failed to generate a valid visualization after multiple attempts.")
+    return None
 
 
 def add_stats_to_figure(fig, df, y_axis, chart_type):
-    # Calculate statistics
+    """
+    Add relevant statistical annotations to the visualization 
+    based on the chart type.
+    """
+    # Check if the y-axis column is numeric
+    if not pd.api.types.is_numeric_dtype(df[y_axis]):
+        st.warning(f"⚠️ Cannot compute statistics for non-numeric column: {y_axis}")
+        return fig
+
+    # Compute statistics for numeric data
     min_val = df[y_axis].min()
     max_val = df[y_axis].max()
     avg_val = df[y_axis].mean()
     median_val = df[y_axis].median()
     std_dev_val = df[y_axis].std()
 
-    # Stats summary text
+    # Format the stats for display
     stats_text = (
         f"📊 **Statistics**\n\n"
         f"- **Min:** ${min_val:,.2f}\n"
@@ -195,97 +261,173 @@ def add_stats_to_figure(fig, df, y_axis, chart_type):
         f"- **Std Dev:** ${std_dev_val:,.2f}"
     )
 
-    # Charts suitable for stats annotations
-    if chart_type in ["bar", "line", "scatter"]:
-        # Add annotation box
+    # Apply stats only to relevant chart types
+    if chart_type in ["bar", "line"]:
+        # Add annotation box for bar and line charts
         fig.add_annotation(
             text=stats_text,
             xref="paper", yref="paper",
-            x=1.05, y=1,
+            x=1.02, y=1,
             showarrow=False,
             align="left",
             font=dict(size=12, color="black"),
-            bordercolor="black",
+            bordercolor="gray",
             borderwidth=1,
-            bgcolor="rgba(255, 255, 255, 0.8)"
+            bgcolor="rgba(255, 255, 255, 0.85)"
         )
 
-        # Add horizontal lines for min, median, avg, max
+        # Add horizontal reference lines
         fig.add_hline(y=min_val, line_dash="dot", line_color="red", annotation_text="Min", annotation_position="bottom right")
         fig.add_hline(y=median_val, line_dash="dash", line_color="orange", annotation_text="Median", annotation_position="top right")
         fig.add_hline(y=avg_val, line_dash="dashdot", line_color="green", annotation_text="Avg", annotation_position="top right")
         fig.add_hline(y=max_val, line_dash="dot", line_color="blue", annotation_text="Max", annotation_position="top right")
 
+    elif chart_type == "scatter":
+        # Add stats annotation only, no lines for scatter plots
+        fig.add_annotation(
+            text=stats_text,
+            xref="paper", yref="paper",
+            x=1.02, y=1,
+            showarrow=False,
+            align="left",
+            font=dict(size=12, color="black"),
+            bordercolor="gray",
+            borderwidth=1,
+            bgcolor="rgba(255, 255, 255, 0.85)"
+        )
+
     elif chart_type == "box":
-        # Box plots already show distribution (no extra stats needed)
+        # Box plots inherently show distribution; no extra stats needed
         pass
 
     elif chart_type == "pie":
-        # Pie charts don't need statistical overlays
-        st.info("📊 Pie charts focus on proportions. No additional stats displayed.")
-    
+        # Pie charts represent proportions, not suitable for stats
+        st.info("📊 Pie charts represent proportions. Additional stats are not applicable.")
+
+    elif chart_type == "heatmap":
+        # Heatmaps already reflect data intensity
+        st.info("📊 Heatmaps inherently reflect distribution. No additional stats added.")
+
     else:
-        st.warning(f"⚠️ No stats added for unsupported chart type: {chart_type}")
+        st.warning(f"⚠️ No statistical overlays applied for unsupported chart type: '{chart_type}'.")
 
     return fig
 
 
-# Function to create PDF with report text and visualizations
-def create_pdf_report_with_viz(report, conclusion, visualizations):
-    pdf = FPDF()
-    pdf.set_auto_page_break(auto=True, margin=15)
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-
-    # Title
-    pdf.set_font("Arial", style="B", size=18)
-    pdf.cell(0, 10, "📊 Analysis Report", ln=True, align="C")
-    pdf.ln(10)
-
-    # Report Content
-    pdf.set_font("Arial", style="B", size=14)
-    pdf.cell(0, 10, "Analysis", ln=True)
-    pdf.set_font("Arial", size=12)
-    pdf.multi_cell(0, 10, report)
-
-    pdf.ln(10)
-    pdf.set_font("Arial", style="B", size=14)
-    pdf.cell(0, 10, "Conclusion", ln=True)
-    pdf.set_font("Arial", size=12)
-    pdf.multi_cell(0, 10, conclusion)
-
-    # Add Visualizations
-    pdf.add_page()
-    pdf.set_font("Arial", style="B", size=16)
-    pdf.cell(0, 10, "📈 Visualizations", ln=True)
-    pdf.ln(5)
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        for i, fig in enumerate(visualizations, start=1):
-            fig_title = fig.layout.title.text if fig.layout.title.text else f"Visualization {i}"
-            x_axis = fig.layout.xaxis.title.text if fig.layout.xaxis.title.text else "X-axis"
-            y_axis = fig.layout.yaxis.title.text if fig.layout.yaxis.title.text else "Y-axis"
-
-            # Save each visualization as a PNG image
-            img_path = os.path.join(temp_dir, f"viz_{i}.png")
-            fig.write_image(img_path)
-
-            # Insert Title and Description
-            pdf.set_font("Arial", style="B", size=14)
-            pdf.multi_cell(0, 10, f"{i}. {fig_title}")
-            pdf.set_font("Arial", size=12)
-            pdf.multi_cell(0, 10, f"X-axis: {x_axis} | Y-axis: {y_axis}")
-            pdf.ln(3)
-
-            # Embed Visualization
-            pdf.image(img_path, w=170)
-            pdf.ln(10)
-
-        # Save PDF
-        temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
-        pdf.output(temp_pdf.name)
-
-    return temp_pdf
+# Dynamically generate Plotly visualizations based on GPT-4o suggestions
+def generate_visualization(suggestion, df):
+    """
+    Generate a Plotly visualization based on GPT-4o's suggestion.
+    If the Y-axis is missing, infer it intelligently.
+    """
+    chart_type = suggestion.get("chart_type", "bar").lower()
+    x_axis = suggestion.get("x_axis")
+    y_axis = suggestion.get("y_axis")
+    group_by = suggestion.get("group_by")
+
+    # Step 1: Infer Y-axis if not provided
+    if not y_axis:
+        numeric_columns = df.select_dtypes(include='number').columns.tolist()
+
+        # Avoid using the same column for both axes
+        if x_axis in numeric_columns:
+            numeric_columns.remove(x_axis)
+
+        # Smart guess: prioritize salary or relevant metrics if available
+        priority_columns = ["salary_in_usd", "income", "earnings", "revenue"]
+        for col in priority_columns:
+            if col in numeric_columns:
+                y_axis = col
+                break
+
+        # Fallback to the first numeric column if no priority columns exist
+        if not y_axis and numeric_columns:
+            y_axis = numeric_columns[0]
+
+    # Step 2: Validate axes
+    if not x_axis or not y_axis:
+        st.warning("⚠️ Unable to determine appropriate columns for visualization.")
+        return None
+
+    #  Step 3: Dynamically select the Plotly function
+    plotly_function = getattr(px, chart_type, None)
+    if not plotly_function:
+        st.warning(f"⚠️ Unsupported chart type '{chart_type}' suggested by GPT-4o.")
+        return None
+
+    #  Step 4: Prepare dynamic plot arguments
+    plot_args = {"data_frame": df, "x": x_axis, "y": y_axis}
+    if group_by and group_by in df.columns:
+        plot_args["color"] = group_by
+
+    try:
+        # Step 5: Generate the visualization
+        fig = plotly_function(**plot_args)
+        fig.update_layout(
+            title=f"{chart_type.title()} Plot of {y_axis.replace('_', ' ').title()} by {x_axis.replace('_', ' ').title()}",
+            xaxis_title=x_axis.replace('_', ' ').title(),
+            yaxis_title=y_axis.replace('_', ' ').title(),
+        )
+
+        # Step 6: Apply statistics intelligently
+        fig = add_statistics_to_visualization(fig, df, y_axis, chart_type)
+
+        return fig
+
+    except Exception as e:
+        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)
+
 
 def escape_markdown(text):
     # Ensure text is a string
@@ -431,31 +573,8 @@ if st.session_state.df is not None:
                 safe_conclusion = escape_markdown(conclusion_result if conclusion_result else "⚠️ No Conclusion Generated.")
                 st.markdown(safe_conclusion)
 
-    # Full Data Visualization Tab
-    with tab2:
-        st.subheader("📊 Comprehensive Data Visualizations")
-
-        fig1 = px.histogram(st.session_state.df, x="job_title", title="Job Title Frequency")
-        st.plotly_chart(fig1)
-
-        fig2 = px.bar(
-            st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
-            x="experience_level", y="salary_in_usd",
-            title="Average Salary by Experience Level"
-        )
-        st.plotly_chart(fig2)
-
-        fig3 = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
-                      title="Salary Distribution by Employment Type")
-        st.plotly_chart(fig3)
-
-    temp_dir.cleanup()
-else:
-    st.info("Please load a dataset to proceed.")
-
 
 # Sidebar Reference
 with st.sidebar:
     st.header("📚 Reference:")
-    st.markdown("[SQL Agents w CrewAI & Llama 3 - Plaban Nayak](https://github.com/plaban1981/Agents/blob/main/SQL_Agents_with_CrewAI_and_Llama_3.ipynb)")
-
+    st.markdown("[SQL Agents w CrewAI & Llama 3 - Plaban Nayak](https://github.com/plaban1981/Agents/blob/main/SQL_Agents_with_CrewAI_and_Llama_3.ipynb)")