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)")