Update agent.py

agent.py

@@ -299,6 +299,7 @@ class AgentState(TypedDict):
     summary: Optional[str]
     interaction_warnings: Optional[List[str]]
     done: Optional[bool]
+    iterations: Optional[int]
 
 # ── Graph Nodes ───────────────────────────────────────────────────────────────
 def agent_node(state: AgentState) -> Dict[str, Any]:
@@ -325,7 +326,7 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
             med = call["args"].get("medication_name", "").lower()
             if not any(
                 c["name"] == "check_drug_interactions" and
-                c["args"].get("potential_prescription","").lower() == med
+                c["args"].get("potential_prescription", "").lower() == med
                 for c in calls
             ):
                 logger.warning(f"Blocking prescribe_medication for '{med}' without interaction check")
@@ -337,7 +338,7 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
             call["args"].setdefault("current_medications", pd.get("medications", {}).get("current", []))
             call["args"].setdefault("allergies", pd.get("allergies", []))
     messages: List[ToolMessage] = []
-    warnings: List[str]     = []
+    warnings: List[str] = []
     try:
         responses = tool_executor.batch(to_execute, return_exceptions=True)
         for call, resp in zip(to_execute, responses):
@@ -388,6 +389,13 @@ 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
+    # Force termination after a set number of iterations to prevent infinite loops.
+    if state["iterations"] >= 4:
+        state["done"] = True
+        return "end_conversation_turn"
     last = state["messages"][-1] if state["messages"] else None
     if not isinstance(last, AIMessage):
         state["done"] = True