Update interim.py

interim.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
@@ -18,18 +17,13 @@ from langchain_community.tools.sql_database.tool import (
     QuerySQLDataBaseTool,
 )
 from langchain_community.utilities.sql_database import SQLDatabase
+from datasets import load_dataset
 import tempfile
 
-# Setup GROQ API Key
+# Setup API key
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-# Callback handler for logging LLM responses
-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
@@ -50,32 +44,44 @@ llm = ChatGroq(
     callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
 )
 
-# App Header
-st.title("SQL-RAG with CrewAI 🚀")
-st.write("Provide your query, and the app will extract, analyze, and summarize the data dynamically.")
-
-# File Upload for Dataset
-uploaded_file = st.file_uploader("Upload your dataset (CSV file)", type=["csv"])
-
-if uploaded_file:
-    st.success("File uploaded successfully!")
-    
-    # Temporary directory for SQLite DB
+st.title("SQL-RAG using CrewAI 🚀")
+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!")
+        except Exception as 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")
-    
-    # Create SQLite database
-    df = pd.read_csv(uploaded_file)
     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)
@@ -90,23 +96,23 @@ if uploaded_file:
 
     # Agents
     sql_dev = Agent(
-        role="Senior Database Developer",
-        goal="Extract data from the database based on user query",
+        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="Senior Data Analyst",
-        goal="Analyze the database response and provide insights",
+        role="Data Analyst",
+        goal="Analyze and provide insights.",
         llm=llm,
         allow_delegation=False,
     )
 
     report_writer = Agent(
-        role="Senior Report Editor",
-        goal="Summarize the analysis into a short report",
+        role="Report Editor",
+        goal="Summarize the analysis.",
         llm=llm,
         allow_delegation=False,
     )
@@ -119,20 +125,19 @@ if uploaded_file:
     )
 
     analyze_data = Task(
-        description="Analyze the data and generate insights for: {query}.",
+        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 concise executive report.",
+        description="Summarize the analysis into a short report.",
         expected_output="Markdown report",
         agent=report_writer,
         context=[analyze_data],
     )
 
-    # Crew
     crew = Crew(
         agents=[sql_dev, data_analyst, report_writer],
         tasks=[extract_data, analyze_data, write_report],
@@ -141,8 +146,7 @@ if uploaded_file:
         memory=False,
     )
 
-    # User Input Query
-    query = st.text_input("Enter your query:")
+    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}
@@ -150,5 +154,6 @@ if uploaded_file:
             st.markdown("### Analysis Report:")
             st.markdown(result)
 
-    # Clean up
-    temp_dir.cleanup()
+    temp_dir.cleanup()
+else:
+    st.warning("Please upload a valid file or provide a correct Hugging Face dataset name.")