Update graph_builder.py

graph_builder.py

@@ -1,19 +1,22 @@
-from langgraph.graph import Graph
+from langgraph.graph import StateGraph
 from llm_node import llm_node
 from formatter_node import formatter_node
 
+class AgentState(dict):
+    pass
+
 def build_graph():
-    graph = Graph()
+    graph = StateGraph(AgentState)
 
     def llm_step(state):
         question = state["question"]
-        result = llm_node().run(question)
-        state["llm_output"] = result
+        llm_output = llm_node(question)
+        state["llm_output"] = llm_output
         return state
 
     def formatter_step(state):
-        cleaned = formatter_node(state["llm_output"])
-        state["final_answer"] = cleaned
+        final_answer = formatter_node(state["llm_output"])
+        state["final_answer"] = final_answer
         return state
 
     graph.add_node("llm", llm_step)