app.py

import streamlit as st
import pandas as pd
import sqlite3
import os
import json
from pathlib import Path
from datetime import datetime, timezone
from crewai import Agent, Crew, Process, Task
from crewai.tools import tool
from langchain_groq import ChatGroq
from langchain.schema.output import LLMResult
from langchain_core.callbacks.base import BaseCallbackHandler
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

# API Key
os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")

# 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=200,
    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
)

st.title("Blah Blah App Using CrewAI 🚀")
st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")

# Initialize session state for data persistence
if "df" not in st.session_state:
    st.session_state.df = None

# 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="HUPD/hupd")
    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.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())