feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'

src/graph.py

@@ -0,0 +1,28 @@
+from typing import Dict, Any
+from langgraph.graph import StateGraph
+from src.nodes import ReflectState, draft_answer, reflect, rewrite
+
+def build_graph() -> StateGraph:
+    graph = StateGraph(ReflectState)
+
+    # Add nodes
+    graph.add_node("draft_answer", draft_answer)
+    graph.add_node("reflect", reflect)
+    graph.add_node("rewrite", rewrite)
+
+    # Define transitions
+    graph.set_entry_point("draft_answer")
+    graph.add_edge("draft_answer", "reflect")
+
+    # Conditional edge after reflect
+    def decide_next(state: ReflectState) -> str:
+        if state["verdict"] == "ok":
+            return "end"
+        if state["round"] < state["max_rounds"]:
+            return "rewrite"
+        return "end"
+
+    graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "end": "end"})
+    graph.add_edge("rewrite", "reflect")
+
+    return graph