Update app.py

app.py

@@ -1,10 +1,7 @@
 import streamlit as st
 import pandas as pd
 import sqlite3
-import tempfile
-from fpdf import FPDF
 import os
-import re
 import json
 from pathlib import Path
 import plotly.express as px
@@ -33,6 +30,7 @@ llm = None
 # Model Selection
 model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
 
+
 # API Key Validation and LLM Initialization
 groq_api_key = os.getenv("GROQ_API_KEY")
 openai_api_key = os.getenv("OPENAI_API_KEY")
@@ -53,12 +51,9 @@ elif model_choice == "GPT-4o":
 # Initialize session state for data persistence
 if "df" not in st.session_state:
     st.session_state.df = None
-if "show_preview" not in st.session_state:
-    st.session_state.show_preview = False  
 
 # Dataset Input
 input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
-
 if input_option == "Use Hugging Face Dataset":
     dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="Einstellung/demo-salaries")
     if st.button("Load Dataset"):
@@ -66,338 +61,16 @@ if input_option == "Use Hugging Face Dataset":
             with st.spinner("Loading dataset..."):
                 dataset = load_dataset(dataset_name, split="train")
                 st.session_state.df = pd.DataFrame(dataset)
-                st.session_state.show_preview = True  # Show preview after loading
                 st.success(f"Dataset '{dataset_name}' loaded successfully!")
+                st.dataframe(st.session_state.df.head())
         except Exception as e:
             st.error(f"Error: {e}")
-
 elif input_option == "Upload CSV File":
     uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
     if uploaded_file:
-        try:
-            st.session_state.df = pd.read_csv(uploaded_file)
-            st.session_state.show_preview = True  # Show preview after loading
-            st.success("File uploaded successfully!")
-        except Exception as e:
-            st.error(f"Error loading file: {e}")
-
-# Show Dataset Preview Only After Loading
-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())
-
-
-
-
-def ask_gpt4o_for_visualization(query, df, llm):
-    columns = ', '.join(df.columns)
-    prompt = f"""
-    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):
-    [
-      {{
-        "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
-
-def add_stats_to_figure(fig, df, y_axis, chart_type):
-    """
-    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()
-
-    # Format the stats for display
-    stats_text = (
-        f"📊 **Statistics**\n\n"
-        f"- **Min:** ${min_val:,.2f}\n"
-        f"- **Max:** ${max_val:,.2f}\n"
-        f"- **Average:** ${avg_val:,.2f}\n"
-        f"- **Median:** ${median_val:,.2f}\n"
-        f"- **Std Dev:** ${std_dev_val:,.2f}"
-    )
-
-    # 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.02, y=1,
-            showarrow=False,
-            align="left",
-            font=dict(size=12, color="black"),
-            bordercolor="gray",
-            borderwidth=1,
-            bgcolor="rgba(255, 255, 255, 0.85)"
-        )
-
-        # 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 inherently show distribution; no extra stats needed
-        pass
-
-    elif chart_type == "pie":
-        # 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 statistical overlays applied for unsupported chart type: '{chart_type}'.")
-
-    return fig
-
-
-# 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)
-
-
-
-# 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)
-    # Escape Markdown characters: *, _, `, ~
-    escape_chars = r"(\*|_|`|~)"
-    return re.sub(escape_chars, r"\\\1", text)
+        st.session_state.df = pd.read_csv(uploaded_file)
+        st.success("File uploaded successfully!")
+        st.dataframe(st.session_state.df.head())
 
 # SQL-RAG Analysis
 if st.session_state.df is not None:
@@ -427,7 +100,6 @@ if st.session_state.df is not None:
         """Validate the SQL query syntax and structure before execution."""
         return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
 
-    # Agents for SQL data extraction and analysis
     sql_dev = Agent(
         role="Senior Database Developer",
         goal="Extract data using optimized SQL queries.",
@@ -445,19 +117,11 @@ if st.session_state.df is not None:
 
     report_writer = Agent(
         role="Technical Report Writer",
-        goal="Write a structured report with Introduction and Key Insights. DO NOT include any Conclusion or Summary.",
-        backstory="Specializes in detailed analytical reports without conclusions.",
+        goal="Summarize the insights into a clear report.",
+        backstory="An expert in summarizing data insights into readable reports.",
         llm=llm,
     )
 
-    conclusion_writer = Agent(
-        role="Conclusion Specialist",
-        goal="Summarize findings into a clear and concise 3-5 line Conclusion highlighting only the most important insights.",
-        backstory="An expert in crafting impactful and clear conclusions.",
-        llm=llm,
-    )
-
-    # Define tasks for report and conclusion
     extract_data = Task(
         description="Extract data based on the query: {query}.",
         expected_output="Database results matching the query.",
@@ -466,100 +130,56 @@ if st.session_state.df is not None:
 
     analyze_data = Task(
         description="Analyze the extracted data for query: {query}.",
-        expected_output="Key Insights and Analysis without any Introduction or Conclusion.",
+        expected_output="Analysis text summarizing findings.",
         agent=data_analyst,
         context=[extract_data],
     )
 
     write_report = Task(
-        description="Write the analysis report with  Introduction and Key Insights. DO NOT include any Conclusion or Summary.",
-        expected_output="Markdown-formatted report excluding Conclusion.",
+        description="Summarize the analysis into an executive report.",
+        expected_output="Markdown report of insights.",
         agent=report_writer,
         context=[analyze_data],
     )
 
-    write_conclusion = Task(
-        description="Summarize the key findings in 3-5 impactful lines, highlighting the maximum, minimum, and average salaries."
-        "Emphasize significant insights on salary distribution and influential compensation trends for strategic decision-making.",
-        expected_output="Markdown-formatted Conclusion section with key insights and statistics.",
-        agent=conclusion_writer,
-        context=[analyze_data],
-    )
-
-
-
-    # Separate Crews for report and conclusion
-    crew_report = Crew(
+    crew = Crew(
         agents=[sql_dev, data_analyst, report_writer],
         tasks=[extract_data, analyze_data, write_report],
         process=Process.sequential,
         verbose=True,
     )
 
-    crew_conclusion = Crew(
-        agents=[data_analyst, conclusion_writer],
-        tasks=[write_conclusion],
-        process=Process.sequential,
-        verbose=True,
-    )
-
-    # Tabs for Query Results and Visualizations
+    # UI: Tabs for Query Results and General Insights
     tab1, tab2 = st.tabs(["🔍 Query Insights + Viz", "📊 Full Data Viz"])
 
-    # Query Insights + Visualization
     with tab1:
         query = st.text_area("Enter Query:", value="Provide insights into the salary of a Principal Data Scientist.")
         if st.button("Submit Query"):
             with st.spinner("Processing query..."):
-                # Step 1: Generate the analysis report
-                report_inputs = {"query": query + " Provide detailed analysis but DO NOT include Conclusion."}
-                report_result = crew_report.kickoff(inputs=report_inputs)
-
-                # Step 2: Generate only the concise conclusion
-                conclusion_inputs = {"query": query + " Provide ONLY the most important insights in 3-5 concise lines."}
-                conclusion_result = crew_conclusion.kickoff(inputs=conclusion_inputs)
-
-                # Step 3: Display the report
-                #st.markdown("### Analysis Report:")
-                st.markdown(report_result if report_result else "⚠️ No Report Generated.")
-
-                # Step 4: Generate Visualizations
-
-
-                # Step 5: Insert Visual Insights
-                st.markdown("### Visual Insights")
+                inputs = {"query": query}
+                result = crew.kickoff(inputs=inputs)
+                st.markdown("### Analysis Report:")
+                st.markdown(result)
 
+                # Query-Specific Visualization
+                if "salary" in query.lower():
+                    fig = px.box(st.session_state.df, x="job_title", y="salary_in_usd", title="Salary Distribution by Job Title")
+                    st.plotly_chart(fig)
 
-                # Step 6: Display Concise Conclusion
-                #st.markdown("#### Conclusion")
-
-                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"
-        )
+        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)")