Update app.py

app.py

@@ -51,24 +51,6 @@ toolkit = SQLDatabaseToolkit(db=db, llm=llm)
 tools = toolkit.get_tools()
 
 
-def create_agent_app(db_path: str):
-    # Construct the SQLite URI from the given file path.
-    db_uri = f"sqlite:///{db_path}"
-    # Create new SQLDatabase connection using the constructed URI.
-    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()
-
-    # ... (rest of the logic remains unchanged)
-    # [Define custom query tool, prompt templates, workflow nodes, etc.]
-
-    # Compile and return the agent application workflow.
-    return workflow.compile()
-    # Define a custom query tool for executing SQL queries
 
 # Define a custom query tool for executing SQL queries
 @tool
@@ -137,78 +119,94 @@ 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"],
+
+def create_agent_app(db_path: str):
+    # Construct the SQLite URI from the given file path.
+    db_uri = f"sqlite:///{db_path}"
+    # Create new SQLDatabase connection using the constructed URI.
+    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()
+
+    # ... (rest of the logic remains unchanged)
+        # Define a custom query tool for executing SQL queries
+
+    # 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()