Update app.py

app.py

@@ -1,17 +1,16 @@
-import json
-import os
+import streamlit as st
+import pandas as pd
 import sqlite3
-from dataclasses import asdict, dataclass
-from datetime import datetime, timezone
+import os
+import json
 from pathlib import Path
-from textwrap import dedent
-from typing import Any, Dict, List, Tuple, Union
-
-import pandas as pd
+from datetime import datetime, timezone
 from crewai import Agent, Crew, Process, Task
 from crewai_tools import tool
-from langchain.schema import AgentFinish
+from langchain_core.prompts import ChatPromptTemplate
+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,
@@ -19,6 +18,137 @@ from langchain_community.tools.sql_database.tool import (
     QuerySQLDataBaseTool,
 )
 from langchain_community.utilities.sql_database import SQLDatabase
-from langchain_core.callbacks.base import BaseCallbackHandler
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_groq import ChatGroq
+import tempfile
+
+# Setup GROQ 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
+
+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 Setup
+llm = ChatGroq(
+    temperature=0,
+    model_name="mixtral-8x7b-32768",
+    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
+)
+
+# App Header
+st.title("Dynamic Query Analysis 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
+    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)
+
+    @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="Senior Database Developer",
+        goal="Extract data from the database based on user query",
+        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",
+        llm=llm,
+        allow_delegation=False,
+    )
+
+    report_writer = Agent(
+        role="Senior Report Editor",
+        goal="Summarize the analysis into a short report",
+        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 and generate insights 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.",
+        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],
+        process=Process.sequential,
+        verbose=2,
+        memory=False,
+    )
+
+    # User Input Query
+    query = st.text_input("Enter your query:")
+    if query:
+        with st.spinner("Processing your query..."):
+            inputs = {"query": query}
+            result = crew.kickoff(inputs=inputs)
+            st.markdown("### Analysis Report:")
+            st.markdown(result)
+
+    # Clean up
+    temp_dir.cleanup()