Update app.py

app.py

@@ -50,6 +50,19 @@ from langchain_community.agent_toolkits import SQLDatabaseToolkit
 toolkit = SQLDatabaseToolkit(db=db, llm=llm)
 tools = toolkit.get_tools()
 
+
+def create_agent_app(db_uri: str):
+    # Create new SQLDatabase connection
+    from langchain_community.utilities import SQLDatabase
+    db_instance = SQLDatabase.from_uri(db_uri)
+
+    # Create SQL toolkit and get the tools
+    from langchain_community.agent_toolkits import SQLDatabaseToolkit
+    toolkit_instance = SQLDatabaseToolkit(db=db_instance, llm=llm)
+    tools_instance = toolkit_instance.get_tools()
+
+    # Define a custom query tool for executing SQL queries
+
 # Define a custom query tool for executing SQL queries
 @tool
 def db_query_tool(query: str) -> str:
@@ -117,111 +130,134 @@ DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the databa
 query_gen_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system), ("placeholder", "{messages}")])
 query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
 
-# Define nodes and fallback functions for the workflow
-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"}])]}
-
-def handle_tool_error(state: State) -> dict:
-    error = state.get("error")
-    tool_calls = state["messages"][-1].tool_calls
-    return {
-        "messages": [
-            ToolMessage(content=f"Error: {repr(error)}\n please fix your mistakes.", tool_call_id=tc["id"])
-            for tc in tool_calls
-        ]
-    }
-
-def create_tool_node_with_fallback(tools_list: list) -> RunnableWithFallbacks[Any, dict]:
-    return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")
-
-def query_gen_node(state: State):
-    message = query_gen.invoke(state)
-    # Check for incorrect tool calls
-    tool_messages = []
-    if message.tool_calls:
-        for tc in message.tool_calls:
-            if tc["name"] != "SubmitFinalAnswer":
-                tool_messages.append(
-                    ToolMessage(
-                        content=f"Error: The wrong tool was called: {tc['name']}. Please fix your mistakes. Remember to only call SubmitFinalAnswer to submit the final answer. Generated queries should be outputted WITHOUT a tool call.",
-                        tool_call_id=tc["id"],
+  # 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"}])]}
+
+    def handle_tool_error(state: State) -> dict:
+        error = state.get("error")
+        tool_calls = state["messages"][-1].tool_calls
+        return {
+            "messages": [
+                ToolMessage(content=f"Error: {repr(error)}\n please fix your mistakes.", tool_call_id=tc["id"])
+                for tc in tool_calls
+            ]
+        }
+
+    def create_tool_node_with_fallback(tools_list: list) -> RunnableWithFallbacks[Any, dict]:
+        return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")
+
+    def query_gen_node(state: State):
+        message = query_gen.invoke(state)
+        # Check for incorrect tool calls
+        tool_messages = []
+        if message.tool_calls:
+            for tc in message.tool_calls:
+                if tc["name"] != "SubmitFinalAnswer":
+                    tool_messages.append(
+                        ToolMessage(
+                            content=f"Error: The wrong tool was called: {tc['name']}. Please fix your mistakes. Remember to only call SubmitFinalAnswer to submit the final answer. Generated queries should be outputted WITHOUT a tool call.",
+                            tool_call_id=tc["id"],
+                        )
                     )
-                )
-    return {"messages": [message] + tool_messages}
-
-def should_continue(state: State) -> Literal[END, "correct_query", "query_gen"]:
-    messages = state["messages"]
-    last_message = messages[-1]
-    if getattr(last_message, "tool_calls", None):
-        return END
-    if last_message.content.startswith("Error:"):
-        return "query_gen"
-    else:
-        return "correct_query"
-
-def model_check_query(state: State) -> dict[str, list[AIMessage]]:
-    """Double-check if the query is correct before executing it."""
-    return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}
-
-# Get tools for listing tables and fetching schema
-list_tables_tool = next((tool for tool in tools if tool.name == "sql_db_list_tables"), None)
-get_schema_tool = next((tool for tool in tools if tool.name == "sql_db_schema"), None)
-
-# Define the workflow (state graph)
-workflow = StateGraph(State)
-workflow.add_node("first_tool_call", first_tool_call)
-workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tables_tool]))
-workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))
-model_get_schema = llm.bind_tools([get_schema_tool])
-workflow.add_node("model_get_schema", lambda state: {"messages": [model_get_schema.invoke(state["messages"])],})
-workflow.add_node("query_gen", query_gen_node)
-workflow.add_node("correct_query", model_check_query)
-workflow.add_node("execute_query", create_tool_node_with_fallback([db_query_tool]))
-
-workflow.add_edge(START, "first_tool_call")
-workflow.add_edge("first_tool_call", "list_tables_tool")
-workflow.add_edge("list_tables_tool", "model_get_schema")
-workflow.add_edge("model_get_schema", "get_schema_tool")
-workflow.add_edge("get_schema_tool", "query_gen")
-workflow.add_conditional_edges("query_gen", should_continue)
-workflow.add_edge("correct_query", "execute_query")
-workflow.add_edge("execute_query", "query_gen")
-
-# Compile the workflow into an agent application.
-agent_app = workflow.compile()
-
-# Initialize Flask and SocketIO
-flask_app = Flask(__name__)
-socketio = SocketIO(flask_app, cors_allowed_origins="*")
-
-# Set up an uploads directory
-UPLOAD_FOLDER = os.path.join(os.getcwd(), "uploads")
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-
-# Create a global agent_app using the default DATABASE_URI
-agent_app = create_agent_app(DATABASE_URI)
-
-# Endpoint for uploading a DB file
-@flask_app.route("/upload", methods=["GET", "POST"])
-def upload():
-    if request.method == "POST":
-        file = request.files.get("file")
-        if not file:
-            return "No file uploaded", 400
-        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
-        file.save(file_path)
-        # Build a new URI (for SQLite, use absolute path)
-        new_db_uri = f"sqlite:///{file_path}"
-        # Reinitialize the agent_app with the new DB
-        global agent_app
-        agent_app = create_agent_app(new_db_uri)
-        socketio.emit("log", {"message": f"[INFO]: Database file '{file.filename}' uploaded and loaded."})
-        return redirect(url_for("index"))
-    return render_template("upload.html")
-
-# Function to run the agent in a separate thread
-def run_agent(prompt):
+        return {"messages": [message] + tool_messages}
+
+    def should_continue(state: State) -> Literal[END, "correct_query", "query_gen"]:
+        messages = state["messages"]
+        last_message = messages[-1]
+        if getattr(last_message, "tool_calls", None):
+            return END
+        if last_message.content.startswith("Error:"):
+            return "query_gen"
+        else:
+            return "correct_query"
+
+    def model_check_query(state: State) -> dict[str, list[AIMessage]]:
+        """Double-check if the query is correct before executing it."""
+        return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}
+
+    # 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)
+
+    # Define the workflow (state graph)
+    workflow = StateGraph(State)
+    workflow.add_node("first_tool_call", first_tool_call)
+    workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tables_tool]))
+    workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))
+    model_get_schema = llm.bind_tools([get_schema_tool])
+    workflow.add_node("model_get_schema", lambda state: {"messages": [model_get_schema.invoke(state["messages"])],})
+    workflow.add_node("query_gen", query_gen_node)
+    workflow.add_node("correct_query", model_check_query)
+    workflow.add_node("execute_query", create_tool_node_with_fallback([db_query_tool]))
+
+    workflow.add_edge(START, "first_tool_call")
+    workflow.add_edge("first_tool_call", "list_tables_tool")
+    workflow.add_edge("list_tables_tool", "model_get_schema")
+    workflow.add_edge("model_get_schema", "get_schema_tool")
+    workflow.add_edge("get_schema_tool", "query_gen")
+    workflow.add_conditional_edges("query_gen", should_continue)
+    workflow.add_edge("correct_query", "execute_query")
+    workflow.add_edge("execute_query", "query_gen")
+
+    # Compile and return the agent application workflow
+    return workflow.compile()
+
+###############################################################################
+# Application Factory: create_app()
+#
+# This function sets up the Flask application, SocketIO, routes, and initializes
+# the global agent_app using the default DATABASE_URI. It returns the Flask app.
+###############################################################################
+def create_app():
+    flask_app = Flask(__name__)
+    socketio = SocketIO(flask_app, cors_allowed_origins="*")
+
+    # Set up an uploads directory (for DB file uploads)
+    UPLOAD_FOLDER = os.path.join(os.getcwd(), "uploads")
+    if not os.path.exists(UPLOAD_FOLDER):
+        os.makedirs(UPLOAD_FOLDER)
+
+    # Create a global agent_app using the default DATABASE_URI
+    global agent_app
+    agent_app = create_agent_app(DATABASE_URI)
+
+    @flask_app.route("/")
+    def index():
+        return render_template("index.html")
+
+    @flask_app.route("/generate", methods=["POST"])
+    def generate():
+        data = request.json
+        prompt = data.get("prompt", "")
+        socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}\n"})
+        # Run the agent in a separate thread
+        thread = threading.Thread(target=run_agent, args=(prompt, socketio))
+        thread.start()
+        return "OK", 200
+
+    @flask_app.route("/upload", methods=["GET", "POST"])
+    def upload():
+        if request.method == "POST":
+            file = request.files.get("file")
+            if not file:
+                return "No file uploaded", 400
+            file_path = os.path.join(UPLOAD_FOLDER, file.filename)
+            file.save(file_path)
+            # For SQLite, use the absolute file path in the URI
+            new_db_uri = f"sqlite:///{file_path}"
+            global agent_app
+            agent_app = create_agent_app(new_db_uri)
+            socketio.emit("log", {"message": f"[INFO]: Database file '{file.filename}' uploaded and loaded."})
+            return redirect(url_for("index"))
+        return render_template("upload.html")
+
+    return flask_app, socketio
+
+###############################################################################
+# Helper function to run the agent; uses the global agent_app.
+###############################################################################
+def run_agent(prompt, socketio):
     try:
         query = {"messages": [("user", prompt)]}
         result = agent_app.invoke(query)
@@ -232,22 +268,8 @@ def run_agent(prompt):
         socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
         socketio.emit("final", {"message": "Generation failed."})
 
-@flask_app.route("/")
-def index():
-    return render_template("index.html")
-
-@flask_app.route("/generate", methods=["POST"])
-def generate():
-    data = request.json
-    prompt = data.get("prompt", "")
-    socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}\n"})
-    # Run the agent in a separate thread
-    thread = threading.Thread(target=run_agent, args=(prompt,))
-    thread.start()
-    return "OK", 200
-
-# Assign the Flask app to "app" for gunicorn
-app = flask_app
+# Create the app and assign to "app" for Gunicorn compatibility.
+app, socketio_instance = create_app()
 
 if __name__ == "__main__":
-    socketio.run(app, debug=True)
+    socketio_instance.run(app, debug=True)