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.")
+st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
 
 # Initialize LLM
 llm = None
@@ -87,152 +87,30 @@ if st.session_state.df is not None and st.session_state.show_preview:
     st.dataframe(st.session_state.df.head())
 
 
-# Helper Function for Validation
-def is_valid_suggestion(suggestion):
-    chart_type = suggestion.get("chart_type", "").lower()
 
-    if chart_type in ["bar", "line", "box", "scatter"]:
-        return all(k in suggestion for k in ["chart_type", "x_axis", "y_axis"])
 
-    elif chart_type == "pie":
-        return all(k in suggestion for k in ["chart_type", "x_axis"])
-
-    elif chart_type == "heatmap":
-        return all(k in suggestion for k in ["chart_type", "x_axis", "y_axis"])
-
-    else:
-        return False
-        
-def ask_gpt4o_for_visualization(query, df, llm, retries=2):
-    import json
-
-    # Identify numeric and categorical columns
-    numeric_columns = df.select_dtypes(include='number').columns.tolist()
-    categorical_columns = df.select_dtypes(exclude='number').columns.tolist()
-
-    # Prompt with Dataset-Specific, Query-Based Examples
+def ask_gpt4o_for_visualization(query, df, llm):
+    columns = ', '.join(df.columns)
     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:
+    Analyze the query and suggest one or more relevant visualizations.
+    Query: "{query}"
+    Available Columns: {columns}
+    Respond in this JSON format (as a list if multiple suggestions):
     [
       {{
-        "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"
+        "chart_type": "bar/box/line/scatter",
+        "x_axis": "column_name",
+        "y_axis": "column_name",
+        "group_by": "optional_column_name"
       }}
     ]
-    **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.
     """
-
-    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}")
-
-        if attempt < retries:
-            st.info("🔄 Retrying visualization suggestion...")
-
-    st.error("❌ Failed to generate a valid visualization after multiple attempts.")
-    return None
-
+    response = llm.generate(prompt)
+    try:
+        return json.loads(response)
+    except json.JSONDecodeError:
+        st.error("⚠️ GPT-4o failed to generate a valid suggestion.")
+        return None
 
 def add_stats_to_figure(fig, df, y_axis, chart_type):
     """
@@ -429,6 +307,91 @@ def handle_visualization_suggestions(suggestions, df):
         st.plotly_chart(fig, use_container_width=True)
 
 
+
+# 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"
+
+    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"
+
+        content += f"\n{i}. {fig_title}\n"
+        content += f"   - X-axis: {x_axis}\n"
+        content += f"   - Y-axis: {y_axis}\n"
+
+        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"
+
+    content += f"\n\n\n{conclusion}"
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode='w', encoding='utf-8') as temp_txt:
+        temp_txt.write(content)
+        return temp_txt.name
+
+
+
+# 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
+
 def escape_markdown(text):
     # Ensure text is a string
     text = str(text)
@@ -573,6 +536,28 @@ 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: