Update app.py

app.py

@@ -240,118 +240,100 @@ def create_agent_app(db_path: str):
 # create_app: The application factory.
 # =============================================================================
     
-def create_app():
-    flask_app = Flask(__name__, static_url_path='/uploads', static_folder='uploads')
-    socketio = SocketIO(flask_app, cors_allowed_origins="*")
+# Flask and SocketIO setup
+flask_app = Flask(__name__)
+flask_app.config['UPLOAD_FOLDER'] = "uploaded_files"
+os.makedirs(flask_app.config['UPLOAD_FOLDER'], exist_ok=True)
 
-    # Ensure uploads folder exists.
-    if not os.path.exists(UPLOAD_FOLDER):
-        os.makedirs(UPLOAD_FOLDER)
-    flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+socketio = SocketIO(flask_app, cors_allowed_origins="*")
 
-    # -------------------------------------------------------------------------
-    # Serve uploaded files via a custom route.
-    # -------------------------------------------------------------------------
-    
-    @flask_app.route("/files/<path:filename>")
-    def uploaded_file(filename):
-        return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
+# Global variables to manage state
+agent_app = None
+db_path = None
+abs_file_path = None
 
-    # -------------------------------------------------------------------------
-    # 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
+# ----------------------------
+# ROUTES
+# ----------------------------
 
-        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."})
+@flask_app.route("/", methods=["GET"])
+def index():
+    return render_template("index.html")
 
-    # -------------------------------------------------------------------------
-    # Route: index page.
-    # -------------------------------------------------------------------------
-    
-    @flask_app.route("/")
-    def index():
-        return render_template("index.html")
+@flask_app.route("/upload", methods=["GET", "POST"])
+def upload():
+    global abs_file_path, db_path, agent_app
+    try:
+        if request.method == "POST":
+            file = request.files.get("file")
+            if not file:
+                return "No file uploaded", 400
 
-    # -------------------------------------------------------------------------
-    # 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
+            filename = secure_filename(file.filename)
+            if filename.endswith('.db'):
+                db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
+                file.save(db_path)
+
+                abs_file_path = os.path.abspath(db_path)
+                agent_app = None  # reset agent so it gets re-initialized on next 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")
+
+
+# ----------------------------
+# AGENT INVOCATION
+# ----------------------------
+
+@socketio.on("user_input")
+def handle_user_input(data):
+    prompt = data.get("message")
+    if not prompt:
+        socketio.emit("log", {"message": "[ERROR]: Empty prompt received."})
+        return
+
+    run_agent(prompt)
+
+
+def run_agent(prompt):
+    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."})
+        return
+
+    try:
+        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)
 
-    # -------------------------------------------------------------------------
-    # 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
         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'):
-                    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")
-
-    return flask_app, socketio
+            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."})
 
-# =============================================================================
-# Create the app for Gunicorn compatibility.
-# =============================================================================
 
-app, socketio_instance = create_app()
+# ----------------------------
+# MAIN
+# ----------------------------
 
 if __name__ == "__main__":
-    socketio_instance.run(app, debug=True)
+    socketio.run(flask_app, debug=True)