Update app.py

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, render_template, request, redirect, url_for, send_from_directory
+from flask import Flask, render_template, request, redirect, url_for, send_from_directory, flash
 from flask_socketio import SocketIO
 import threading
 import os
@@ -62,6 +62,15 @@ def create_agent_app(db_path: str):
     # -------------------------------------------------------------------------
     @tool
     def db_query_tool(query: str) -> str:
+        """
+        Executes a SQL query on the connected SQLite database.
+
+        Parameters:
+            query (str): A SQL query string to be executed.
+
+        Returns:
+            str: The result from the database if successful, or an error message if not.
+        """
         result = db_instance.run_no_throw(query)
         return result if result else "Error: Query failed. Please rewrite your query and try again."
 
@@ -116,7 +125,7 @@ def create_agent_app(db_path: str):
         "If you get an empty result set, you should try to rewrite the query to get a non-empty result set.\n"
         "NEVER make stuff up if you don't have enough information to answer the query... just say you don't have enough information.\n\n"
         "If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.\n"
-        "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. Do not return any sql query except answer."
+        "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. Do not return any SQL query except answer."
     )
     query_gen_prompt = ChatPromptTemplate.from_messages([
         ("system", query_gen_system),
@@ -124,25 +133,32 @@ def create_agent_app(db_path: str):
     ])
     query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
 
-    # Update database URI and file path
+    # -------------------------------------------------------------------------
+    # Update database URI and file path, 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}"
     print("db_uri", db_uri)
-    flash(f"db_uri:{db_uri}", "warning")
-    # Create SQLDatabase connection using langchain utility.
+    # Uncomment if flash is needed; ensure you have flask.flash imported if so.
+    # flash(f"db_uri:{db_uri}", "warning")
+
     from langchain_community.utilities import SQLDatabase
     db_instance = SQLDatabase.from_uri(db_uri)
     print("db_instance----->", db_instance)
-    flash(f"db_instance:{db_instance}", "warning")
+    # flash(f"db_instance:{db_instance}", "warning")
 
+    # -------------------------------------------------------------------------
     # Create SQL toolkit.
+    # -------------------------------------------------------------------------
     from langchain_community.agent_toolkits import SQLDatabaseToolkit
     toolkit_instance = SQLDatabaseToolkit(db=db_instance, llm=llm)
     tools_instance = toolkit_instance.get_tools()
 
+    # -------------------------------------------------------------------------
     # Define workflow nodes and fallback functions.
+    # -------------------------------------------------------------------------
     def first_tool_call(state: State) -> dict[str, list[AIMessage]]:
         return {"messages": [AIMessage(content="", tool_calls=[{"name": "sql_db_list_tables", "args": {}, "id": "tool_abcd123"}])]}
 
@@ -181,7 +197,9 @@ def create_agent_app(db_path: str):
     def model_check_query(state: State) -> dict[str, list[AIMessage]]:
         return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}
 
-    # Get table listing and schema tools.
+    # -------------------------------------------------------------------------
+    # Get tools for listing tables and fetching schema.
+    # -------------------------------------------------------------------------
     list_tables_tool = next((tool for tool in tools_instance if tool.name == "sql_db_list_tables"), None)
     get_schema_tool = next((tool for tool in tools_instance if tool.name == "sql_db_schema"), None)
 
@@ -207,11 +225,11 @@ def create_agent_app(db_path: str):
     # Return compiled workflow
     return workflow.compile()
 
+
 # =============================================================================
 # create_app: The application factory.
 # =============================================================================
 def create_app():
-    # Configure static folder for uploads.
     flask_app = Flask(__name__, static_url_path='/uploads', static_folder='uploads')
     socketio = SocketIO(flask_app, cors_allowed_origins="*")
 
@@ -220,7 +238,9 @@ def create_app():
         os.makedirs(UPLOAD_FOLDER)
     flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 
+    # -------------------------------------------------------------------------
     # 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)
@@ -229,26 +249,26 @@ def create_app():
     # Helper: run_agent runs the agent with the given prompt.
     # -------------------------------------------------------------------------
     def run_agent(prompt, socketio):
-        global agent_app, abs_file_path
+        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 init
+            # 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:
@@ -256,16 +276,15 @@ def create_app():
             socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
             socketio.emit("final", {"message": "Generation failed."})
 
-
     # -------------------------------------------------------------------------
-    # Route: index page
+    # Route: index page.
     # -------------------------------------------------------------------------
     @flask_app.route("/")
     def index():
         return render_template("index.html")
 
     # -------------------------------------------------------------------------
-    # Route: generate (POST) – receives a prompt, runs the agent.
+    # Route: generate (POST) – receives a prompt and runs the agent.
     # -------------------------------------------------------------------------
     @flask_app.route("/generate", methods=["POST"])
     def generate():
@@ -280,7 +299,6 @@ def create_app():
             return "OK", 200
         except Exception as e:
             print(f"[ERROR]: {str(e)}")
-            flash(f"[ERROR]: {str(e)}", "error")
             socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
             return "ERROR", 500
 
@@ -289,17 +307,19 @@ def create_app():
     # -------------------------------------------------------------------------
     @flask_app.route("/upload", methods=["GET", "POST"])
     def upload():
-        global abs_file_path, db_path
+        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, don't create agent
+                    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"))