From 1f8ee7b070f31bce022314da8954335e7cff8570 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?=
 =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= <mirzagitov03@mail.ru>
Date: Thu, 18 Jun 2026 10:08:05 +0000
Subject: [PATCH] Solution ready for publish: update main.py

---
 main.py | 40 +++++++++++++++++-----------------------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/main.py b/main.py
index dcfc8eb..f3a229e 100644
--- a/main.py
+++ b/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,29 +40,21 @@ 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. "
-            "Respond in JSON with keys 'verdict' and 'critique'."
-        )
-        response = llm.invoke(prompt)
-        # Expect JSON; simple parse
-        import json
-        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}"
+    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. "
+        "Respond in JSON with keys 'verdict' and 'critique'."
+    )
+    response = llm.invoke(prompt)
+    # Expect JSON; simple parse
+    import json
+    data = json.loads(response.content)
+    verdict = data.get("verdict", "needs_revision")
+    critique = data.get("critique", "")
     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()