Update src/app/graph/workflow.py

src/app/graph/workflow.py

@@ -1,52 +1,54 @@
-from langgraph.graph import StateGraph
-from typing import TypedDict
-from app.agents.analyzer import analyzer_node
-from app.agents.generator import generator_node
-from app.agents.reviewer import reviewer_node
-from app.agents.refactor import refactor_node
-from app.config import MAX_ITERATIONS
-
-
-class CodeState(TypedDict):
-    prompt: str
-    structured_prompt: str
-    code: str
-    review: str
-    iteration: int
-
-
-def should_continue(state):
-    if state["iteration"] >= MAX_ITERATIONS:
-        return "end"
-
-    if "critical" in state["review"].lower():
-        return "refactor"
-
-    return "end"
-
-
-def build_graph():
-    builder = StateGraph(CodeState)
-
-    builder.add_node("analyzer", analyzer_node)
-    builder.add_node("generator", generator_node)
-    builder.add_node("reviewer", reviewer_node)
-    builder.add_node("refactor", refactor_node)
-
-    builder.set_entry_point("analyzer")
-
-    builder.add_edge("analyzer", "generator")
-    builder.add_edge("generator", "reviewer")
-
-    builder.add_conditional_edges(
-        "reviewer",
-        should_continue,
-        {
-            "refactor": "refactor",
-            "end": None
-        }
-    )
-
-    builder.add_edge("refactor", "reviewer")
-
+from langgraph.graph import StateGraph
+from typing import TypedDict
+from app.agents.analyzer import analyzer_node
+from app.agents.generator import generator_node
+from app.agents.reviewer import reviewer_node
+from app.agents.refactor import refactor_node
+from app.config import MAX_ITERATIONS
+from langgraph.graph import END
+
+class CodeState(TypedDict):
+    prompt: str
+    structured_prompt: str
+    code: str
+    review: str
+    iteration: int
+
+
+def should_continue(state):
+    review = state.get("review", "").lower()
+
+    if state["iteration"] >= MAX_ITERATIONS:
+        return "end"
+
+    if any(word in review for word in ["critical", "error", "bug"]):
+        return "refactor"
+
+    return "end"
+
+
+def build_graph():
+    builder = StateGraph(CodeState)
+
+    builder.add_node("analyzer", analyzer_node)
+    builder.add_node("generator", generator_node)
+    builder.add_node("reviewer", reviewer_node)
+    builder.add_node("refactor", refactor_node)
+
+    builder.set_entry_point("analyzer")
+
+    builder.add_edge("analyzer", "generator")
+    builder.add_edge("generator", "reviewer")
+
+    builder.add_conditional_edges(
+        "reviewer",
+        should_continue,
+        {
+            "refactor": "refactor",
+            "end": END
+        }
+    )
+
+    builder.add_edge("refactor", "reviewer")
+
     return builder.compile()