Update app.py

app.py

@@ -7,7 +7,6 @@ from pathlib import Path
 from datetime import datetime, timezone
 from crewai import Agent, Crew, Process, Task
 from crewai_tools import tool
-from langchain_core.prompts import ChatPromptTemplate
 from langchain_groq import ChatGroq
 from langchain.schema.output import LLMResult
 from langchain_core.callbacks.base import BaseCallbackHandler
@@ -21,14 +20,10 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
+# Setup API key
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-class Event:
-    def __init__(self, event, text):
-        self.event = event
-        self.timestamp = datetime.now(timezone.utc).isoformat()
-        self.text = text
-
+# Callback handler for logging
 class LLMCallbackHandler(BaseCallbackHandler):
     def __init__(self, log_path: Path):
         self.log_path = log_path
@@ -42,6 +37,7 @@ class LLMCallbackHandler(BaseCallbackHandler):
         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 Setup
 llm = ChatGroq(
     temperature=0,
     model_name="mixtral-8x7b-32768",
@@ -49,102 +45,115 @@ llm = ChatGroq(
 )
 
 st.title("SQL-RAG using CrewAI 🚀")
-st.write("Analyze and summarize Hugging Face datasets using natural language queries with SQL-based retrieval.")
-
-default_dataset = "datascience/ds-salaries"
-st.text("Example dataset: `datascience/ds-salaries` (You can enter your own dataset name)")
-
-dataset_name = st.text_input("Enter Hugging Face dataset name:", value=default_dataset)
-
-if dataset_name:
-    with st.spinner("Loading dataset..."):
+st.write("Analyze and summarize data using natural language queries with SQL-based retrieval.")
+
+# File upload or Hugging Face dataset input
+option = st.radio("Choose your input method:", ["Upload a CSV file", "Enter Hugging Face dataset name"])
+
+if option == "Upload a CSV file":
+    uploaded_file = st.file_uploader("Upload your dataset (CSV format)", type=["csv"])
+    if uploaded_file:
+        df = pd.read_csv(uploaded_file)
+        st.success("File uploaded successfully!")
+else:
+    dataset_name = st.text_input("Enter Hugging Face dataset name:", placeholder="e.g., imdb, ag_news")
+    if dataset_name:
         try:
             dataset = load_dataset(dataset_name, split="train")
             df = pd.DataFrame(dataset)
             st.success(f"Dataset '{dataset_name}' loaded successfully!")
-            st.write("Preview of the dataset:")
-            st.dataframe(df.head())
-
-            temp_dir = tempfile.TemporaryDirectory()
-            db_path = os.path.join(temp_dir.name, "data.db")
-            connection = sqlite3.connect(db_path)
-            df.to_sql("data_table", connection, if_exists="replace", index=False)
-            db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
-
-            @tool("list_tables")
-            def list_tables() -> str:
-                return ListSQLDatabaseTool(db=db).invoke("")
-
-            @tool("tables_schema")
-            def tables_schema(tables: str) -> str:
-                return InfoSQLDatabaseTool(db=db).invoke(tables)
-
-            @tool("execute_sql")
-            def execute_sql(sql_query: str) -> str:
-                return QuerySQLDataBaseTool(db=db).invoke(sql_query)
-
-            @tool("check_sql")
-            def check_sql(sql_query: str) -> str:
-                return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
-
-            sql_dev = Agent(
-                role="Database Developer",
-                goal="Extract data from the database.",
-                llm=llm,
-                tools=[list_tables, tables_schema, execute_sql, check_sql],
-                allow_delegation=False,
-            )
-
-            data_analyst = Agent(
-                role="Data Analyst",
-                goal="Analyze and provide insights.",
-                llm=llm,
-                allow_delegation=False,
-            )
-
-            report_writer = Agent(
-                role="Report Editor",
-                goal="Summarize the analysis.",
-                llm=llm,
-                allow_delegation=False,
-            )
-
-            extract_data = Task(
-                description="Extract data required for the query: {query}.",
-                expected_output="Database result for the query",
-                agent=sql_dev,
-            )
-
-            analyze_data = Task(
-                description="Analyze the data for: {query}.",
-                expected_output="Detailed analysis text",
-                agent=data_analyst,
-                context=[extract_data],
-            )
-
-            write_report = Task(
-                description="Summarize the analysis into a short report.",
-                expected_output="Markdown report",
-                agent=report_writer,
-                context=[analyze_data],
-            )
-
-            crew = Crew(
-                agents=[sql_dev, data_analyst, report_writer],
-                tasks=[extract_data, analyze_data, write_report],
-                process=Process.sequential,
-                verbose=2,
-                memory=False,
-            )
-
-            query = st.text_input("Enter your query:", placeholder="e.g., 'How does salary vary by company size?'")
-            if query:
-                with st.spinner("Processing your query..."):
-                    inputs = {"query": query}
-                    result = crew.kickoff(inputs=inputs)
-                    st.markdown("### Analysis Report:")
-                    st.markdown(result)
-
-            temp_dir.cleanup()
         except Exception as e:
-            st.error(f"Error loading dataset: {e}")
+            st.error(f"Error loading Hugging Face dataset: {e}")
+            df = None
+
+if 'df' in locals() and not df.empty:
+    st.write("### Dataset Preview:")
+    st.dataframe(df.head())
+
+    # Create a temporary SQLite database
+    temp_dir = tempfile.TemporaryDirectory()
+    db_path = os.path.join(temp_dir.name, "data.db")
+    connection = sqlite3.connect(db_path)
+    df.to_sql("data_table", connection, if_exists="replace", index=False)
+    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
+
+    # Tools
+    @tool("list_tables")
+    def list_tables() -> str:
+        return ListSQLDatabaseTool(db=db).invoke("")
+
+    @tool("tables_schema")
+    def tables_schema(tables: str) -> str:
+        return InfoSQLDatabaseTool(db=db).invoke(tables)
+
+    @tool("execute_sql")
+    def execute_sql(sql_query: str) -> str:
+        return QuerySQLDataBaseTool(db=db).invoke(sql_query)
+
+    @tool("check_sql")
+    def check_sql(sql_query: str) -> str:
+        return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
+
+    # Agents
+    sql_dev = Agent(
+        role="Database Developer",
+        goal="Extract data from the database.",
+        llm=llm,
+        tools=[list_tables, tables_schema, execute_sql, check_sql],
+        allow_delegation=False,
+    )
+
+    data_analyst = Agent(
+        role="Data Analyst",
+        goal="Analyze and provide insights.",
+        llm=llm,
+        allow_delegation=False,
+    )
+
+    report_writer = Agent(
+        role="Report Editor",
+        goal="Summarize the analysis.",
+        llm=llm,
+        allow_delegation=False,
+    )
+
+    # Tasks
+    extract_data = Task(
+        description="Extract data required for the query: {query}.",
+        expected_output="Database result for the query",
+        agent=sql_dev,
+    )
+
+    analyze_data = Task(
+        description="Analyze the data for: {query}.",
+        expected_output="Detailed analysis text",
+        agent=data_analyst,
+        context=[extract_data],
+    )
+
+    write_report = Task(
+        description="Summarize the analysis into a short report.",
+        expected_output="Markdown report",
+        agent=report_writer,
+        context=[analyze_data],
+    )
+
+    crew = Crew(
+        agents=[sql_dev, data_analyst, report_writer],
+        tasks=[extract_data, analyze_data, write_report],
+        process=Process.sequential,
+        verbose=2,
+        memory=False,
+    )
+
+    query = st.text_input("Enter your query:", placeholder="e.g., 'What are the top 5 highest salaries?'")
+    if query:
+        with st.spinner("Processing your query..."):
+            inputs = {"query": query}
+            result = crew.kickoff(inputs=inputs)
+            st.markdown("### Analysis Report:")
+            st.markdown(result)
+
+    temp_dir.cleanup()
+else:
+    st.warning("Please upload a valid file or provide a correct Hugging Face dataset name.")