Updated LangGraph reflection agent with max_rounds logic.: update main.py

main.py

@@ -67,7 +67,12 @@ builder.set_entry_point("draft_answer")
 builder.add_edge("draft_answer", "reflect")
 builder.add_conditional_edges(
     "reflect",
-    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
+    lambda state: (
+        "END"
+        if state["verdict"] == "ok"
+        or state.get("round", 0) >= state.get("max_rounds", 2)
+        else "rewrite"
+    ),
 )
 builder.add_edge("rewrite", "reflect")