Update app.py

app.py

@@ -131,29 +131,88 @@ def create_agent_app(db_path: str):
 
     return workflow.compile()
 
-@flask_app.route("/", methods=["GET"])
+@flask_app.route("/files/<path:filename>")
+def uploaded_file(filename):
+    return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
+# -------------------------------------------------------------------------
+# Helper: run_agent runs the agent with the given prompt.
+# -------------------------------------------------------------------------
+
+def run_agent(prompt, socketio):
+    global agent_app, abs_file_path, db_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."})
+        return
+    try:
+        # Lazy agent initialization: use the previously uploaded DB.
+        if agent_app is None:
+            print("[INFO]: Initializing agent for the first time...")
+            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."
+        print("final_answer------>", result)
+        socketio.emit("final", {"message": result})
+    except Exception as e:
+        print(f"[ERROR]: {str(e)}")
+        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+        socketio.emit("final", {"message": "Generation failed."})
+# -------------------------------------------------------------------------
+# 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..."})
+        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()
+        return "OK", 200
+    except Exception as e:
+        print(f"[ERROR]: {str(e)}")
+        socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+        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
+    global abs_file_path, agent_app, db_path
     try:
         if request.method == "POST":
             file = request.files.get("file")
             if not file:
+                print("No file uploaded")
                 return "No file uploaded", 400
-
             filename = secure_filename(file.filename)
             if filename.endswith('.db'):
-                save_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
-                file.save(save_path)
-                abs_file_path = os.path.abspath(save_path)
-                agent_app = None  # Reset agent; reinitialize on next query.
-                socketio.emit("log", {"message": f"Database '{filename}' uploaded."})
+                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"))
         return render_template("upload.html")
     except Exception as e:
+        print(f"[ERROR]: {str(e)}")
         socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
         return render_template("upload.html")
 
@@ -165,6 +224,7 @@ def handle_user_input(data):
         return
     run_agent(prompt)
 
+'''
 def run_agent(prompt):
     global agent_app, abs_file_path
     if not abs_file_path:
@@ -185,6 +245,7 @@ def run_agent(prompt):
     except Exception as e:
         socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
         socketio.emit("final", {"message": "Generation failed."})
+'''
 
 # Expose the Flask app as "app" for Gunicorn
 app = flask_app