Update app.py

app.py

@@ -39,8 +39,6 @@ load_dotenv()
 UPLOAD_FOLDER = os.path.join(os.getcwd(), "uploads")
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
-DATABASE_URI = f"sqlite:///{os.path.join(BASE_DIR, 'data', 'mydb.db')}"
-print("DATABASE URI:", DATABASE_URI)
 
 # API Keys from .env file
 os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
@@ -130,12 +128,9 @@ def create_agent_app(db_path: str):
     query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
 
     # -------------------------------------------------------------------------
-    # Update database URI, create SQLDatabase connection.
+    # Create SQLDatabase connection.
     # -------------------------------------------------------------------------
-    abs_db_path_local = os.path.abspath(db_path)
-    global DATABASE_URI
-    DATABASE_URI = abs_db_path_local
-    db_uri = f"sqlite:///{abs_db_path_local}"
+    db_uri = f"sqlite:///{os.path.abspath(db_path)}"
     print("db_uri", db_uri)
 
     try:
@@ -248,16 +243,33 @@ def create_app():
         try:
             if agent_app is None:
                 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")
+                try:
+                    agent_app = create_agent_app(abs_file_path)
+                    socketio.emit("log", {"message": "[INFO]: Agent initialized."})
+                    flash("Agent initialized.", "info")
+                except Exception as e:
+                    error_message = f"Agent initialization failed: {str(e)}"
+                    socketio.emit("log", {"message": f"[ERROR]: {error_message}"})
+                    socketio.emit("final", {"message": "Agent initialization failed."})
+                    flash(error_message, "error")
+                    return
+
             query = {"messages": [("user", prompt)]}
             result = agent_app.invoke(query)
             try:
-                result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
+                # Attempt to extract the final answer from the tool call arguments
+                if result and result["messages"] and len(result["messages"]) > 0 and result["messages"][-1].tool_calls and len(result["messages"][-1].tool_calls) > 0:
+                    result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
+                else:
+                    result = "Query execution did not return a valid final answer."
+                    flash("Query execution did not return a valid final answer.", "warning")
+            except KeyError:
+                result = "The agent's response did not contain the expected 'final_answer' key."
+                flash("Unexpected agent response format.", "warning")
             except Exception as e:
-                result = "Query failed or no valid answer found."
-                flash("Query failed or no valid answer found.", "warning")
+                result = f"An error occurred while processing the agent's response: {str(e)}"
+                flash("Error processing agent response.", "error")
+
             print("final_answer------>", result)
             socketio.emit("final", {"message": result})
         except Exception as e:
@@ -289,6 +301,14 @@ def create_app():
             flash(error_message, "error")
             return "ERROR", 500
 
+    def is_sqlite_db(file_path):
+        try:
+            with open(file_path, 'rb') as f:
+                header = f.read(16)
+                return header[:16] == b'SQLite format 3\x00'
+        except Exception:
+            return False
+
     @flask_app.route("/upload", methods=["GET", "POST"])
     def upload():
         global abs_file_path, agent_app, db_path
@@ -303,12 +323,20 @@ def create_app():
                     db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
                     print("Saving file to:", db_path)
                     file.save(db_path)
+                    if not is_sqlite_db(db_path):
+                        os.remove(db_path)
+                        flash("Uploaded file is not a valid SQLite database.", "error")
+                        socketio.emit("log", {"message": "[ERROR]: Invalid database file uploaded."})
+                        return render_template("upload.html")
                     abs_file_path = os.path.abspath(db_path)
-                    agent_app = None   # Reset the agent to be lazily reinitialized 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."})
                     flash(f"Database file '{filename}' uploaded successfully.", "info")
+                    agent_app = None   # Reset the agent to be lazily reinitialized on next query.
                     return redirect(url_for("index"))
+                else:
+                    flash("Invalid file format. Please upload a .db file.", "error")
+                    return render_template("upload.html")
             return render_template("upload.html")
         except Exception as e:
             error_message = f"[ERROR]: {str(e)}"