Update interim.py

interim.py

@@ -4,12 +4,13 @@ import sqlite3
 import os
 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 crewai.tools import tool
 from langchain_groq import ChatGroq
+from langchain_openai import ChatOpenAI
 from langchain.schema.output import LLMResult
-from langchain_core.callbacks.base import BaseCallbackHandler
 from langchain_community.tools.sql_database.tool import (
     InfoSQLDatabaseTool,
     ListSQLDatabaseTool,
@@ -20,39 +21,41 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
-# API Key
-os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
+st.title("SQL-RAG Using CrewAI 🚀")
+st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
 
 # Initialize LLM
-class LLMCallbackHandler(BaseCallbackHandler):
-    def __init__(self, log_path: Path):
-        self.log_path = log_path
-
-    def on_llm_start(self, serialized, prompts, **kwargs):
-        with self.log_path.open("a", encoding="utf-8") as file:
-            file.write(json.dumps({"event": "llm_start", "text": prompts[0], "timestamp": datetime.now().isoformat()}) + "\n")
-
-    def on_llm_end(self, response: LLMResult, **kwargs):
-        generation = response.generations[-1][-1].message.content
-        with self.log_path.open("a", encoding="utf-8") as file:
-            file.write(json.dumps({"event": "llm_end", "text": generation, "timestamp": datetime.now().isoformat()}) + "\n")
-
-llm = ChatGroq(
-    temperature=0,
-    model_name="groq/llama-3.3-70b-versatile",
-    max_tokens=1024,
-    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
-)
+llm = None
 
-st.title("SQL-RAG Using CrewAI 🚀")
-st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
+# 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"):
@@ -60,16 +63,25 @@ 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:
-        st.session_state.df = pd.read_csv(uploaded_file)
-        st.success("File uploaded successfully!")
-        st.dataframe(st.session_state.df.head())
+        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())
 
 # SQL-RAG Analysis
 if st.session_state.df is not None:
@@ -86,19 +98,20 @@ if st.session_state.df is not None:
 
     @tool("tables_schema")
     def tables_schema(tables: str) -> str:
-        """Get schema and sample rows for given tables."""
+        """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."""
+        """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:
-        """Check the validity of a SQL query."""
+        """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.",
@@ -116,11 +129,19 @@ if st.session_state.df is not None:
 
     report_writer = Agent(
         role="Technical Report Writer",
-        goal="Summarize the insights into a clear report.",
-        backstory="An expert in summarizing data insights into readable reports.",
+        goal="Write a structured report with Key Insights and Analysis. DO NOT include Introduction or Conclusion.",
+        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.",
@@ -129,33 +150,111 @@ if st.session_state.df is not None:
 
     analyze_data = Task(
         description="Analyze the extracted data for query: {query}.",
-        expected_output="Analysis text summarizing findings.",
+        expected_output="Key Insights and Analysis without any Introduction or Conclusion.",
         agent=data_analyst,
         context=[extract_data],
     )
 
     write_report = Task(
-        description="Summarize the analysis into an executive report.",
-        expected_output="Markdown report of insights.",
+        description="Write the analysis report with Key Insights. DO NOT include a Conclusion.",
+        expected_output="Markdown-formatted report excluding Conclusion.",
         agent=report_writer,
         context=[analyze_data],
     )
 
-    crew = Crew(
+    write_conclusion = Task(
+        description="Write a brief and impactful 3-5 line Conclusion summarizing only the most important insights.",
+        expected_output="Markdown-formatted concise Conclusion section.",
+        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,
     )
 
-    query = st.text_area("Enter Query:", placeholder="e.g., 'What is the average salary for senior employees?'")
-    if st.button("Submit Query"):
-        with st.spinner("Processing query..."):
-            inputs = {"query": query}
-            result = crew.kickoff(inputs=inputs)
-            st.markdown("### Analysis Report:")
-            st.markdown(result)
+    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")
+                st.markdown(conclusion_result if conclusion_result else "⚠️ No Conclusion Generated.")
+
+    # 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.")
+    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)")