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 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") # 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()) #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 # 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) # 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) # 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 # Step 5: Insert Visual Insights st.markdown("### Visual Insights") # 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)")