Update agent.py

agent.py

@@ -154,7 +154,7 @@ def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
             flags.append(f"Red Flag: Hypertensive urgency/emergency ({sys}/{dia} mmHg).")
         if sys <= 90 or dia <= 60:
             flags.append(f"Red Flag: Hypotension ({sys}/{dia} mmHg).")
-    return list(dict.fromkeys(flags))  # dedupe, preserve order
+    return list(dict.fromkeys(flags))
 
 def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
     """Format patient_data dict into a markdown-like prompt section."""
@@ -301,6 +301,13 @@ class AgentState(TypedDict):
     done: Optional[bool]
     iterations: Optional[int]
 
+# Helper to propagate state fields
+def propagate_state(new: Dict[str, Any], old: Dict[str, Any]) -> Dict[str, Any]:
+    for key in ["iterations", "done", "patient_data", "summary", "interaction_warnings"]:
+        if key in old and key not in new:
+            new[key] = old[key]
+    return new
+
 # ── Graph Nodes ───────────────────────────────────────────────────────────────
 def agent_node(state: AgentState) -> Dict[str, Any]:
     msgs = state["messages"]
@@ -309,16 +316,19 @@ def agent_node(state: AgentState) -> Dict[str, Any]:
     logger.info(f"Invoking LLM with {len(msgs)} messages")
     try:
         response = model_with_tools.invoke(msgs)
-        return {"messages": [response]}
+        new_state = {"messages": [response]}
+        return propagate_state(new_state, state)
     except Exception as e:
         logger.exception("Error in agent_node")
-        return {"messages": [AIMessage(content=f"Error: {e}")]} 
+        new_state = {"messages": [AIMessage(content=f"Error: {e}")]}
+        return propagate_state(new_state, state)
 
 def tool_node(state: AgentState) -> Dict[str, Any]:
     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")
-        return {"messages": [], "interaction_warnings": None}
+        new_state = {"messages": []}
+        return propagate_state(new_state, state)
     calls = last.tool_calls
     blocked_ids = set()
     for call in calls:
@@ -360,20 +370,23 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
                 tool_call_id=call["id"],
                 name=call["name"]
             ))
-    return {"messages": messages, "interaction_warnings": warnings or None}
+    new_state = {"messages": messages, "interaction_warnings": warnings or None}
+    return propagate_state(new_state, state)
 
 def reflection_node(state: AgentState) -> Dict[str, Any]:
     warns = state.get("interaction_warnings")
     if not warns:
         logger.warning("reflection_node called without warnings")
-        return {"messages": [], "interaction_warnings": None}
+        new_state = {"messages": []}
+        return propagate_state(new_state, state)
     triggering = None
     for msg in reversed(state["messages"]):
         if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
             triggering = msg
             break
     if not triggering:
-        return {"messages": [AIMessage(content="Internal Error: reflection context missing.")], "interaction_warnings": None}
+        new_state = {"messages": [AIMessage(content="Internal Error: reflection context missing.")]}
+        return propagate_state(new_state, state)
     prompt = (
         "You are SynapseAI, performing a focused safety review of the following plan:\n\n"
         f"{triggering.content}\n\n"
@@ -382,17 +395,20 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
     )
     try:
         resp = llm.invoke([SystemMessage(content="Safety reflection"), HumanMessage(content=prompt)])
-        return {"messages": [AIMessage(content=resp.content)], "interaction_warnings": None}
+        new_state = {"messages": [AIMessage(content=resp.content)]}
+        return propagate_state(new_state, state)
     except Exception as e:
         logger.exception("Error during reflection")
-        return {"messages": [AIMessage(content=f"Error during reflection: {e}")], "interaction_warnings": None}
+        new_state = {"messages": [AIMessage(content=f"Error during reflection: {e}")]}
+        return propagate_state(new_state, state)
 
 # ── Routing Functions ────────────────────────────────────────────────────────
 def should_continue(state: AgentState) -> str:
     # Initialize or increment the iteration counter
     state.setdefault("iterations", 0)
     state["iterations"] += 1
-    # Force termination after a set number of iterations to prevent infinite loops.
+    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"