Create super_flwed_dynamic_viz_v2.py

mylab/super_flwed_dynamic_viz_v2.py

@@ -0,0 +1,614 @@
+import streamlit as st
+import pandas as pd
+import sqlite3
+import tempfile
+from fpdf import FPDF
+import threading
+import time
+import os
+import re
+import json
+from pathlib import Path
+import plotly.express as px
+from datetime import datetime, timezone
+from crewai import Agent, Crew, Process, Task
+from crewai.tools import tool
+from langchain_groq import ChatGroq
+from langchain_openai import ChatOpenAI
+from langchain.schema.output import LLMResult
+from langchain_community.tools.sql_database.tool import (
+    InfoSQLDatabaseTool,
+    ListSQLDatabaseTool,
+    QuerySQLCheckerTool,
+    QuerySQLDataBaseTool,
+)
+from langchain_community.utilities.sql_database import SQLDatabase
+from datasets import load_dataset
+import tempfile
+
+st.title("SQL-RAG Using CrewAI 🚀")
+st.write("Analyze datasets using natural language queries.")
+
+# Initialize LLM
+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")
+
+if model_choice == "llama-3.3-70b":
+    if not groq_api_key:
+        st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
+        llm = None
+    else:
+        llm = ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
+elif model_choice == "GPT-4o":
+    if not openai_api_key:
+        st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
+        llm = None
+    else:
+        llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
+
+if llm is None:
+    st.error("❌ LLM is not initialized. Please check your API keys and model selection.")
+
+# 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"):
+        try:
+            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!")
+        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())
+
+
+# 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
+    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.
+    """
+
+    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
+
+
+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)
+
+
+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)
+
+
+# SQL-RAG Analysis
+if st.session_state.df is not None:
+    temp_dir = tempfile.TemporaryDirectory()
+    db_path = os.path.join(temp_dir.name, "data.db")
+    connection = sqlite3.connect(db_path)
+    st.session_state.df.to_sql("salaries", connection, if_exists="replace", index=False)
+    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
+
+    @tool("list_tables")
+    def list_tables() -> str:
+        """List all tables in the database."""
+        return ListSQLDatabaseTool(db=db).invoke("")
+
+    @tool("tables_schema")
+    def tables_schema(tables: str) -> str:
+        """Get the schema and sample rows for the specified tables."""
+        return InfoSQLDatabaseTool(db=db).invoke(tables)
+
+    @tool("execute_sql")
+    def execute_sql(sql_query: str) -> str:
+        """Execute a SQL query against the database and return the results."""
+        return QuerySQLDataBaseTool(db=db).invoke(sql_query)
+
+    @tool("check_sql")
+    def check_sql(sql_query: str) -> str:
+        """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.",
+        backstory="An expert in writing optimized SQL queries for complex databases.",
+        llm=llm,
+        tools=[list_tables, tables_schema, execute_sql, check_sql],
+    )
+
+    data_analyst = Agent(
+        role="Senior Data Analyst",
+        goal="Analyze the data and produce insights.",
+        backstory="A seasoned analyst who identifies trends and patterns in datasets.",
+        llm=llm,
+    )
+
+    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.",
+        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.",
+        agent=sql_dev,
+    )
+
+    analyze_data = Task(
+        description="Analyze the extracted data for query: {query}.",
+        expected_output="Key Insights and Analysis without any Introduction or Conclusion.",
+        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.",
+        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(
+        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
+    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"):
+            result_container = {"report": None, "conclusion": None, "visuals": None}
+            progress_bar = st.progress(0, text="🚀 Starting Analysis...")
+
+            # Define parallel tasks
+            def generate_report():
+                progress_bar.progress(20, text="📝 Generating Analysis Report...")
+                report_inputs = {"query": query + " Provide detailed analysis but DO NOT include Conclusion."}
+                result_container['report'] = crew_report.kickoff(inputs=report_inputs)
+                progress_bar.progress(40, text="✅ Analysis Report Ready!")
+
+            def generate_conclusion():
+                progress_bar.progress(40, text="📝 Crafting Conclusion...")
+                conclusion_inputs = {"query": query + " Provide ONLY the most important insights in 3-5 concise lines."}
+                result_container['conclusion'] = crew_conclusion.kickoff(inputs=conclusion_inputs)
+                progress_bar.progress(60, text="✅ Conclusion Ready!")
+
+            def generate_visuals():
+                progress_bar.progress(60, text="📊 Creating Visualizations...")
+                result_container['visuals'] = ask_gpt4o_for_visualization(query, st.session_state.df, llm)
+                progress_bar.progress(80, text="✅ Visualizations Ready!")
+
+            # Run tasks in parallel
+            thread_report = threading.Thread(target=generate_report)
+            thread_conclusion = threading.Thread(target=generate_conclusion)
+            thread_visuals = threading.Thread(target=generate_visuals)
+
+            thread_report.start()
+            thread_conclusion.start()
+            thread_visuals.start()
+
+            # Wait for all threads to finish
+            thread_report.join()
+            thread_conclusion.join()
+            thread_visuals.join()
+
+            progress_bar.progress(100, text="✅ Full Analysis Complete!")
+            time.sleep(0.5)
+            progress_bar.empty()
+
+            # Display Report
+            st.markdown("## 📊 Analysis Report")
+            st.markdown(result_container['report'] if result_container['report'] else "⚠️ No Report Generated.")
+
+            # Display Visual Insights
+            st.markdown("## 📈 Visual Insights")
+            if result_container['visuals']:
+                handle_visualization_suggestions(result_container['visuals'], st.session_state.df)
+            else:
+                st.warning("⚠️ No suitable visualizations to display.")
+
+            # Display Conclusion
+            st.markdown("## 📝 Conclusion")
+            safe_conclusion = escape_markdown(result_container['conclusion'] if result_container['conclusion'] else "⚠️ No Conclusion Generated.")
+            st.markdown(safe_conclusion)
+
+
+# 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)")