Update app.py

app.py

@@ -137,78 +137,78 @@ 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 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"],
-                        )
+# 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_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()
+                )
+    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()