Update agent.py

agent.py

@@ -310,6 +310,9 @@ def propagate_state(new: Dict[str, Any], old: Dict[str, Any]) -> Dict[str, Any]:
 
 # ── Graph Nodes ───────────────────────────────────────────────────────────────
 def agent_node(state: AgentState) -> Dict[str, Any]:
+    # Check for termination
+    if state.get("done", False):
+        return state
     msgs = state["messages"]
     if not msgs or not isinstance(msgs[0], SystemMessage):
         msgs = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + msgs
@@ -324,6 +327,8 @@ def agent_node(state: AgentState) -> Dict[str, Any]:
         return propagate_state(new_state, state)
 
 def tool_node(state: AgentState) -> Dict[str, Any]:
+    if state.get("done", False):
+        return state
     last = state["messages"][-1]
     if not isinstance(last, AIMessage) or not getattr(last, "tool_calls", None):
         logger.warning("tool_node invoked without pending tool_calls")
@@ -374,6 +379,8 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
     return propagate_state(new_state, state)
 
 def reflection_node(state: AgentState) -> Dict[str, Any]:
+    if state.get("done", False):
+        return state
     warns = state.get("interaction_warnings")
     if not warns:
         logger.warning("reflection_node called without warnings")
@@ -404,11 +411,9 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
 
 # ── Routing Functions ────────────────────────────────────────────────────────
 def should_continue(state: AgentState) -> str:
-    # Initialize or increment the iteration counter
     state.setdefault("iterations", 0)
     state["iterations"] += 1
     logger.info(f"Iteration count: {state['iterations']}")
-    # Force termination after a set number of iterations
     if state["iterations"] >= 4:
         state["done"] = True
         return "end_conversation_turn"