Create dummy_funcs.py

mylab/dummy_funcs.py

@@ -0,0 +1,489 @@
+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
+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
+from difflib import get_close_matches
+import tempfile
+
+st.title("SQL-RAG Using CrewAI 🚀")
+st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
+
+# 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")
+
+# 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())
+
+# 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)
+
+
+# Synonym mapping for flexible query understanding
+COLUMN_SYNONYMS = {
+    "job_title": ["job title", "job role", "role", "designation", "position", "job responsibility", "occupation"],
+    "experience_level": ["experience level", "seniority", "experience", "career stage", "years of experience"],
+    "employment_type": ["employment type", "job type", "contract type", "employment status"],
+    "salary_in_usd": ["salary", "income", "earnings", "pay", "wage", "compensation"],
+    "remote_ratio": ["remote work", "work from home", "remote ratio", "remote", "telecommute"],
+    "company_size": ["company size", "organization size", "business size", "firm size"],
+    "employee_residence": ["country", "residence", "location", "employee location", "home country"],
+    "company_location": ["company location", "office location", "company country", "headquarters", "location", "located", "area"],
+}
+
+
+# Fuzzy matcher for mapping query terms to dataset columns
+def fuzzy_match_columns(query):
+    query = query.lower()
+    all_synonyms = {synonym: col for col, synonyms in COLUMN_SYNONYMS.items() for synonym in synonyms}
+    words = query.replace("and", "").replace("vs", "").replace("by", "").split()
+
+    matched_columns = []
+    for word in words:
+        matches = get_close_matches(word, all_synonyms.keys(), n=1, cutoff=0.6)
+        matched_columns.extend([all_synonyms[match] for match in matches])
+
+    return list(dict.fromkeys(matched_columns))
+
+# Ask LLM to suggest relevant columns if fuzzy matching fails
+def ask_llm_for_columns(query, llm, df):
+    columns = ', '.join(df.columns)
+    prompt = f"""
+    Analyze this user query and suggest the most relevant dataset columns for visualization.
+    Query: "{query}"
+    Available Columns: {columns}
+    Respond in this JSON format:
+    {{
+      "x_axis": "column_name",
+      "y_axis": "column_name",
+      "group_by": "optional_column_name"
+    }}
+    """
+
+    response = llm.generate(prompt)
+    try:
+        suggestion = json.loads(response)
+        return suggestion
+    except json.JSONDecodeError:
+        st.error("⚠️ Failed to interpret AI response. Please refine your query.")
+        return None
+
+# Add min, max, and average salary annotations to the chart
+def add_stats_to_figure(fig, df, y_axis):
+    min_salary = df[y_axis].min()
+    max_salary = df[y_axis].max()
+    avg_salary = df[y_axis].mean()
+
+    fig.add_annotation(
+        text=f"Min: ${min_salary:,.2f} | Max: ${max_salary:,.2f} | Avg: ${avg_salary:,.2f}",
+        xref="paper", yref="paper",
+        x=0.5, y=1.1,
+        showarrow=False,
+        font=dict(size=12, color="black"),
+        bgcolor="rgba(255, 255, 255, 0.7)"
+    )
+    return fig
+
+# Unified Visualization Generator with Fuzzy Matching and LLM Fallback
+def generate_visual_from_query(query, df, llm=None):
+    try:
+        # Step 1: Attempt Fuzzy Matching
+        matched_columns = fuzzy_match_columns(query)
+
+        # Step 2: Fallback to LLM if no columns are matched
+        if not matched_columns and llm:
+            st.info("🤖 No match found. Asking AI for suggestions...")
+            suggestion = ask_llm_for_columns(query, llm, df)
+            if suggestion:
+                matched_columns = [suggestion.get("x_axis"), suggestion.get("group_by")]
+
+        # Step 3: Process Matched Columns
+        if len(matched_columns) >= 2:
+            x_axis, group_by = matched_columns[0], matched_columns[1]
+        elif len(matched_columns) == 1:
+            x_axis, group_by = matched_columns[0], None
+        else:
+            st.warning("❓ No matching columns found. Try rephrasing your query.")
+            return None
+
+        # Step 4: Visualization Generation
+
+        # Distribution Plot
+        if "distribution" in query:
+            fig = px.box(df, x=x_axis, y="salary_in_usd", color=group_by,
+                         title=f"Salary Distribution by {x_axis.replace('_', ' ').title()}"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return add_stats_to_figure(fig, df, "salary_in_usd")
+
+        # Average Salary Plot
+        elif "average" in query or "mean" in query:
+            grouped_df = df.groupby([x_axis] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
+            fig = px.bar(grouped_df, x=x_axis, y="salary_in_usd", color=group_by,
+                         title=f"Average Salary by {x_axis.replace('_', ' ').title()}"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return add_stats_to_figure(fig, df, "salary_in_usd")
+
+        # Salary Trends Over Time
+        elif "trend" in query and "work_year" in df.columns:
+            grouped_df = df.groupby(["work_year", x_axis])["salary_in_usd"].mean().reset_index()
+            fig = px.line(grouped_df, x="work_year", y="salary_in_usd", color=x_axis,
+                          title=f"Salary Trend Over Years by {x_axis.replace('_', ' ').title()}")
+            return add_stats_to_figure(fig, df, "salary_in_usd")
+
+        # Remote Work Impact
+        elif "remote" in query:
+            grouped_df = df.groupby(["remote_ratio"] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
+            fig = px.bar(grouped_df, x="remote_ratio", y="salary_in_usd", color=group_by,
+                         title="Remote Work Impact on Salary")
+            return add_stats_to_figure(fig, df, "salary_in_usd")
+
+        # No Specific Match
+        else:
+            st.warning("⚠️ No suitable visualization to display!")
+            return None
+
+    except Exception as e:
+        st.error(f"Error generating visualization: {e}")
+        return None
+
+
+# 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"):
+            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
+                visualizations = []
+
+                fig_salary = px.box(st.session_state.df, x="job_title", y="salary_in_usd",
+                                    title="Salary Distribution by Job Title")
+                visualizations.append(fig_salary)
+
+                fig_experience = 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"
+                )
+                visualizations.append(fig_experience)
+
+                fig_employment = px.box(st.session_state.df, x="employment_type", y="salary_in_usd",
+                                        title="Salary Distribution by Employment Type")
+                visualizations.append(fig_employment)
+
+                # Step 5: Insert Visual Insights
+                st.markdown("### Visual Insights")
+                for fig in visualizations:
+                    st.plotly_chart(fig, use_container_width=True)
+
+                # 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"
+        )
+        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)")
+