Update agent.py

agent.py

@@ -871,11 +871,7 @@ def call_llm(state):
     response = llm.invoke(messages)
     return {"messages": messages + [response]}
 
-builder.set_entry_point("call_llm")
-builder.add_conditional_edges("call_llm", should_call_tool, {
-    "call_tool": "call_tool",
-    "end": None
-})
+
 
 # After tool runs, go back to the LLM
 builder.add_edge("call_tool", "call_llm")
@@ -1137,7 +1133,10 @@ def build_graph() -> StateGraph:
     call_llm = llm.bind_tools(tools)
 
     # --- Build the graph ---
-    builder = StateGraph(AgentState)
+from langgraph.graph import StateGraph
+
+def build_graph():
+    builder = StateGraph(AgentState)  # Keep your typed state graph
 
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
@@ -1145,12 +1144,15 @@ def build_graph() -> StateGraph:
     builder.add_node("call_tool", tool_dispatcher)
 
     builder.set_entry_point("retriever")
+
     builder.add_edge("retriever", "assistant")
     builder.add_edge("assistant", "call_llm")
+
     builder.add_conditional_edges("call_llm", should_call_tool, {
         "call_tool": "call_tool",
         "end": None
     })
+
     builder.add_edge("call_tool", "call_llm")
 
     graph = builder.compile()