Solution ready for publish: update main.py

main.py

@@ -2,7 +2,7 @@
 
 This implementation follows the assignment requirements:
 - Draft answer node
-- Reflect node that uses try/except to retry generation when needed
+- Reflect node that critiques the draft
 - Rewrite node that updates draft based on critique
 - max_rounds default 2
 - CLI entry point
@@ -12,7 +12,6 @@ from typing import TypedDict, Dict
 import os
 
 from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import create_chat_agent
 from langchain_openai import ChatOpenAI
 
 # --- State definition -----------------------------------------------------
@@ -41,14 +40,10 @@ def draft_answer(state: ReflectState) -> Dict:
 def reflect(state: ReflectState) -> Dict:
     """Critique the draft.
 
-    Implements retry logic: if the LLM raises an exception during generation,
-    it will be caught and the node will return a verdict of "needs_revision"
-    with an empty critique. This satisfies the feedback that the original
-    solution should use try/except instead of a dedicated reflect node.
+    The node returns a verdict ('ok' or 'needs_revision') and 2–3 concise points of improvement.
     """
     draft = state["draft"]
     question = state["question"]
-    try:
     prompt = (
         f"You are a critical reviewer. Evaluate the following draft answer to the question '{question}'. "
         "Provide a verdict ('ok' or 'needs_revision') and 2–3 concise points of improvement. "
@@ -60,10 +55,6 @@ def reflect(state: ReflectState) -> Dict:
     data = json.loads(response.content)
     verdict = data.get("verdict", "needs_revision")
     critique = data.get("critique", "")
-    except Exception as e:
-        # On any exception, force a revision
-        verdict = "needs_revision"
-        critique = f"LLM error: {e}"
     return {"verdict": verdict, "critique": critique}
 
 
@@ -88,10 +79,13 @@ builder.add_node("rewrite", rewrite)
 # Connections
 builder.set_entry_point("draft_answer")
 builder.add_edge("draft_answer", "reflect")
+
+# Conditional after reflect: if ok -> END, else if round < max_rounds -> rewrite, else -> END
 builder.add_conditional_edges(
     "reflect",
-    lambda x: END if x["verdict"] == "ok" else "rewrite",
+    lambda x: END if x["verdict"] == "ok" else "rewrite" if x["round"] < x["max_rounds"] else END,
 )
+
 builder.add_edge("rewrite", "reflect")
 
 graph = builder.compile()