Update graph_builder.py

graph_builder.py

@@ -1,21 +1,38 @@
 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):
-    pass
+    """
+    Holds the agent state passed between graph nodes.
+    Always includes:
+    - question: str
+    - llm_output: str (added after LLM node)
+    - final_answer: str (added after formatter)
+    """
+    def __init__(self, initial_data):
+        super().__init__()
+        if "question" not in initial_data:
+            raise ValueError("AgentState requires a 'question' key at init")
+        self.update(initial_data)
 
 def build_graph():
     graph = StateGraph(AgentState)
 
-    def llm_step(state):
-        question = state["question"]
+    def llm_step(state, config: RunnableConfig = None):
+        question = state.get("question")
+        if not question:
+            raise ValueError("Missing 'question' in state")
         llm_output = llm_node(question)
         state["llm_output"] = llm_output
         return state
 
-    def formatter_step(state):
-        final_answer = formatter_node(state["llm_output"])
+    def formatter_step(state, config: RunnableConfig = None):
+        llm_output = state.get("llm_output")
+        if not llm_output:
+            raise ValueError("Missing 'llm_output' in state")
+        final_answer = formatter_node(llm_output)
         state["final_answer"] = final_answer
         return state