Update graph_builder.py

graph_builder.py

@@ -1,7 +1,4 @@
 from langgraph.graph import StateGraph
-from langchain_core.runnables import RunnableConfig
-from llm_node import llm_node
-from formatter_node import formatter_node
 
 class AgentState(dict):
     """
@@ -17,10 +14,13 @@ class AgentState(dict):
             raise ValueError("AgentState requires a 'question' key at init")
         self.update(initial_data)
 
+from llm_node import llm_node
+from formatter_node import formatter_node
+
 def build_graph():
     graph = StateGraph(AgentState)
 
-    def llm_step(state, config: RunnableConfig = None):
+    def llm_step(state):
         question = state.get("question")
         if not question:
             raise ValueError("Missing 'question' in state")
@@ -28,7 +28,7 @@ def build_graph():
         state["llm_output"] = llm_output
         return state
 
-    def formatter_step(state, config: RunnableConfig = None):
+    def formatter_step(state):
         llm_output = state.get("llm_output")
         if not llm_output:
             raise ValueError("Missing 'llm_output' in state")
@@ -38,7 +38,7 @@ def build_graph():
 
     graph.add_node("llm", llm_step)
     graph.add_node("formatter", formatter_step)
-    graph.set_entry_point("llm", input_key="question")
+    graph.set_entry_point("llm")  # ✅ no input_key
     graph.add_edge("llm", "formatter")
     graph.set_finish_point("formatter")