Update app.py

app.py

@@ -12,10 +12,9 @@ from langchain_openai import ChatOpenAI
 from langchain_community.agent_toolkits import create_sql_agent
 from huggingface_hub import InferenceClient
 from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import ToolExecutor, ToolInvocation
+from langgraph.prebuilt import ToolNode
 from langchain_core.tools import tool
 from langgraph.checkpoint.sqlite import SqliteSaver
-from langchain.agents import tool
 import gradio as gr
 
 load_dotenv()
@@ -178,28 +177,28 @@ def refine_answer(answer: str, question: str) -> str:
     )
     return response["choices"][0]["message"]["content"]
 
-tool_executor = ToolExecutor(tools=[llama_to_sql, execute_sql_query, refine_answer])
+tool_node = ToolNode([llama_to_sql, execute_sql_query, refine_answer])
 
 class GraphState(dict):
     question: str
-    sql_query: str
-    sql_response: str
-    final_answer: str
+    tool: str
+    tool_input: dict
+    result: str
 
 async def node_generate_sql(state: GraphState):
-    result = await tool_executor.ainvoke({"question": state["question"]})
-    state["sql_query"] = result["output"]
-    return state
+    state["tool"] = "llama_to_sql"
+    state["tool_input"] = {"question": state["question"]}
+    return await tool_node.ainvoke(state)
 
 async def node_execute_sql(state: GraphState):
-    result = await tool_executor.ainvoke({"sql_query": state["sql_query"]})
-    state["sql_response"] = result["output"]
-    return state
+    state["tool"] = "execute_sql_query"
+    state["tool_input"] = {"sql_query": state["result"]}
+    return await tool_node.ainvoke(state)
 
 async def node_refine_if_enabled(state: GraphState):
-    result = await tool_executor.ainvoke({"answer": state["sql_response"], "question": state["question"]})
-    state["final_answer"] = result["output"]
-    return state
+    state["tool"] = "refine_answer"
+    state["tool_input"] = {"answer": state["result"], "question": state["question"]}
+    return await tool_node.ainvoke(state)
 
 builder = StateGraph(GraphState)
 builder.add_node("generate_sql", node_generate_sql)
@@ -214,9 +213,9 @@ memory = SqliteSaver("graph_checkpoint")
 graph = builder.compile(checkpointer=memory)
 
 async def async_chatbot_response(question, model_name):
-    inputs = {"question": question, "model": model_name}
+    inputs = {"question": question}
     result = await graph.ainvoke(inputs)
-    return result["final_answer"]
+    return result["result"]
 
 def toggle_history():
     global show_history_flag