Update app.py

app.py

@@ -1,6 +1,7 @@
-from flask import Flask, render_template, request, redirect, url_for, flash
+from flask import Flask, render_template, request, redirect, url_for, flash, send_from_directory
 from flask_socketio import SocketIO
 import os
+import threading
 from dotenv import load_dotenv
 from werkzeug.utils import secure_filename
 
@@ -34,6 +35,8 @@ os.environ["MISTRAL_API_KEY"] = os.getenv("MISTRAL_API_KEY")
 # Flask and SocketIO setup
 flask_app = Flask(__name__)
 flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+# Set secret key for flash messages:
+flask_app.config['SECRET_KEY'] = os.getenv("FLASK_SECRET_KEY", "mysecretkey")
 socketio = SocketIO(flask_app, cors_allowed_origins="*")
 
 # Global state
@@ -41,16 +44,28 @@ agent_app = None
 abs_file_path = None
 
 def create_agent_app(db_path: str):
-    from langchain_groq import ChatGroq
-    llm = ChatGroq(model="llama3-70b-8192")
+    try:
+        from langchain_groq import ChatGroq
+        llm = ChatGroq(model="llama3-70b-8192")
+    except Exception as e:
+        flash(f"[ERROR]: Failed to initialize ChatGroq: {e}", "error")
+        raise
 
     abs_db_path = os.path.abspath(db_path)
-    db_instance = SQLDatabase.from_uri(f"sqlite:///{abs_db_path}")
+    try:
+        db_instance = SQLDatabase.from_uri(f"sqlite:///{abs_db_path}")
+    except Exception as e:
+        flash(f"[ERROR]: Failed to connect to DB: {e}", "error")
+        raise
 
     @tool
     def db_query_tool(query: str) -> str:
-        result = db_instance.run_no_throw(query)
-        return result or "Error: Query failed. Please rewrite your query and try again."
+        try:
+            result = db_instance.run_no_throw(query)
+            return result or "Error: Query failed. Please rewrite your query and try again."
+        except Exception as e:
+            flash(f"[ERROR]: Exception during query execution: {e}", "error")
+            return f"Error: {str(e)}"
 
     class SubmitFinalAnswer(BaseModel):
         final_answer: str = Field(...)
@@ -58,18 +73,26 @@ def create_agent_app(db_path: str):
     class State(TypedDict):
         messages: Annotated[list[AnyMessage], add_messages]
 
-    query_check = ChatPromptTemplate.from_messages([
-        ("system", "You are a SQL expert. Fix common issues in SQLite queries."),
-        ("placeholder", "{messages}")
-    ]) | llm.bind_tools([db_query_tool])
-
-    query_gen = ChatPromptTemplate.from_messages([
-        ("system", "You are a SQL expert. Generate SQLite query and return answer using SubmitFinalAnswer tool."),
-        ("placeholder", "{messages}")
-    ]) | llm.bind_tools([SubmitFinalAnswer])
-
-    toolkit = SQLDatabaseToolkit(db=db_instance, llm=llm)
-    tools_instance = toolkit.get_tools()
+    try:
+        query_check = ChatPromptTemplate.from_messages([
+            ("system", "You are a SQL expert. Fix common issues in SQLite queries."),
+            ("placeholder", "{messages}")
+        ]) | llm.bind_tools([db_query_tool])
+    
+        query_gen = ChatPromptTemplate.from_messages([
+            ("system", "You are a SQL expert. Generate SQLite query and return answer using SubmitFinalAnswer tool."),
+            ("placeholder", "{messages}")
+        ]) | llm.bind_tools([SubmitFinalAnswer])
+    except Exception as e:
+        flash(f"[ERROR]: Failed to create prompt templates: {e}", "error")
+        raise
+    
+    try:
+        toolkit = SQLDatabaseToolkit(db=db_instance, llm=llm)
+        tools_instance = toolkit.get_tools()
+    except Exception as e:
+        flash(f"[ERROR]: Failed to initialize SQL toolkit: {e}", "error")
+        raise
 
     def first_tool_call(state: State):
         return {"messages": [AIMessage(content="", tool_calls=[{"name": "sql_db_list_tables", "args": {}, "id": "tool_abcd123"}])]} 
@@ -84,7 +107,11 @@ def create_agent_app(db_path: str):
         return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")
 
     def query_gen_node(state: State):
-        message = query_gen.invoke(state)
+        try:
+            message = query_gen.invoke(state)
+        except Exception as e:
+            flash(f"[ERROR]: Exception in query_gen_node: {e}", "error")
+            raise
         tool_messages = []
         if message.tool_calls:
             for tc in message.tool_calls:
@@ -114,7 +141,7 @@ def create_agent_app(db_path: str):
     workflow.add_node("first_tool_call", first_tool_call)
     workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tool]))
     workflow.add_node("get_schema_tool", create_tool_node_with_fallback([schema_tool]))
-    # Corrected the unterminated string literal in the lambda function below:
+    # Fixed unterminated string literal:
     workflow.add_node("model_get_schema", lambda s: {"messages": [model_get_schema.invoke(s["messages"])]})
     workflow.add_node("query_gen", query_gen_node)
     workflow.add_node("correct_query", model_check_query)
@@ -133,87 +160,106 @@ def create_agent_app(db_path: str):
 
 @flask_app.route("/files/<path:filename>")
 def uploaded_file(filename):
-    return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
+    try:
+        return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
+    except Exception as e:
+        flash(f"[ERROR]: Could not send file: {str(e)}", "error")
+        return redirect(url_for("index"))
+
 # -------------------------------------------------------------------------
 # Helper: run_agent runs the agent with the given prompt.
 # -------------------------------------------------------------------------
-
 def run_agent(prompt, socketio):
-    global agent_app, abs_file_path, db_path
+    global agent_app, abs_file_path
     if not abs_file_path:
         socketio.emit("log", {"message": "[ERROR]: No DB file uploaded."})
         socketio.emit("final", {"message": "No database available. Please upload one and try again."})
+        flash("No database available. Please upload one and try again.", "error")
         return
     try:
         # Lazy agent initialization: use the previously uploaded DB.
         if agent_app is None:
-            print("[INFO]: Initializing agent for the first time...")
+            socketio.emit("log", {"message": "[INFO]: Initializing agent for the first time..."})
             agent_app = create_agent_app(abs_file_path)
             socketio.emit("log", {"message": "[INFO]: Agent initialized."})
+            flash("Agent initialized.", "info")
         query = {"messages": [("user", prompt)]}
         result = agent_app.invoke(query)
         try:
             result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
-        except Exception:
+        except Exception as e:
             result = "Query failed or no valid answer found."
-        print("final_answer------>", result)
+            flash("Query failed or no valid answer found.", "warning")
         socketio.emit("final", {"message": result})
     except Exception as e:
-        print(f"[ERROR]: {str(e)}")
-        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+        error_message = f"Generation failed: {str(e)}"
+        socketio.emit("log", {"message": f"[ERROR]: {error_message}"})
         socketio.emit("final", {"message": "Generation failed."})
+        flash(error_message, "error")
+        traceback.print_exc()
+
 # -------------------------------------------------------------------------
 # Route: index page.
 # -------------------------------------------------------------------------
-
 @flask_app.route("/")
 def index():
     return render_template("index.html")
+
 # -------------------------------------------------------------------------
 # Route: generate (POST) – receives a prompt and runs the agent.
 # -------------------------------------------------------------------------
-
 @flask_app.route("/generate", methods=["POST"])
 def generate():
     try:
-        socketio.emit("log", {"message": "[STEP]: Entering query_gen..."})
+        socketio.emit("log", {"message": "[STEP]: Entering query generation..."})
         data = request.json
         prompt = data.get("prompt", "")
         socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}"})
         thread = threading.Thread(target=run_agent, args=(prompt, socketio))
         socketio.emit("log", {"message": f"[INFO]: Starting thread: {thread}"})
         thread.start()
+        flash("Query submitted successfully.", "info")
         return "OK", 200
     except Exception as e:
-        print(f"[ERROR]: {str(e)}")
-        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+        error_message = f"[ERROR]: {str(e)}"
+        socketio.emit("log", {"message": error_message})
+        flash(error_message, "error")
         return "ERROR", 500
+
 # -------------------------------------------------------------------------
 # Route: upload (GET/POST) – handles uploading the SQLite DB file.
 # -------------------------------------------------------------------------
-
 @flask_app.route("/upload", methods=["GET", "POST"])
 def upload():
-    global abs_file_path, agent_app, db_path
+    global abs_file_path, agent_app
     try:
         if request.method == "POST":
             file = request.files.get("file")
             if not file:
-                print("No file uploaded")
+                flash("No file uploaded.", "error")
                 return "No file uploaded", 400
             filename = secure_filename(file.filename)
             if filename.endswith('.db'):
                 db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
-                print("Saving file to:", db_path)
-                file.save(db_path)
-                abs_file_path = os.path.abspath(db_path)  # Save it here; agent init will occur on first query.
-                print(f"[INFO]: File '{filename}' uploaded. Agent will be initialized on first query.")
-                socketio.emit("log", {"message": f"[INFO]: Database file '{filename}' uploaded."})
-                return redirect(url_for("index"))
+                try:
+                    file.save(db_path)
+                    abs_file_path = os.path.abspath(db_path)  # Save it here; agent init will occur on first query.
+                    agent_app = None  # Reset agent on upload.
+                    flash(f"Database file '{filename}' uploaded successfully.", "info")
+                    socketio.emit("log", {"message": f"[INFO]: Database file '{filename}' uploaded."})
+                    return redirect(url_for("index"))
+                except Exception as save_err:
+                    flash(f"Error saving file: {save_err}", "error")
+                    socketio.emit("log", {"message": f"[ERROR]: Error saving file: {save_err}"})
+                    return render_template("upload.html")
+            else:
+                flash("Only .db files are allowed.", "error")
+                return render_template("upload.html")
         return render_template("upload.html")
     except Exception as e:
-        print(f"[ERROR]: {str(e)}")
-        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+        error_message = f"[ERROR]: {str(e)}"
+        flash(error_message, "error")
+        socketio.emit("log", {"message": error_message})
         return render_template("upload.html")
 
 @socketio.on("user_input")
@@ -221,31 +267,9 @@ def handle_user_input(data):
     prompt = data.get("message")
     if not prompt:
         socketio.emit("log", {"message": "[ERROR]: Empty prompt."})
+        flash("Empty prompt.", "error")
         return
-    run_agent(prompt)
-
-'''
-def run_agent(prompt):
-    global agent_app, abs_file_path
-    if not abs_file_path:
-        socketio.emit("final", {"message": "No DB uploaded."})
-        return
-    try:
-        if agent_app is None:
-            agent_app = create_agent_app(abs_file_path)
-            socketio.emit("log", {"message": "[INFO]: Agent initialized."})
-
-        query = {"messages": [("user", prompt)]}
-        result = agent_app.invoke(query)
-        try:
-            result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
-        except Exception:
-            result = "Query failed or no valid answer found."
-        socketio.emit("final", {"message": result})
-    except Exception as e:
-        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
-        socketio.emit("final", {"message": "Generation failed."})
-'''
+    run_agent(prompt, socketio)
 
 # Expose the Flask app as "app" for Gunicorn
 app = flask_app