Update src/graph.py

src/graph.py

@@ -21,13 +21,13 @@ llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
 agent = create_agent(model=llm, tools=[])
 
 # --- Node implementations -----------------------------------------------
-async def draft_answer(state: CodeReviewState) -> Dict[str, str]:
+async def draft_review(state: CodeReviewState) -> Dict[str, str]:
     code = state["code"]
     prompt = (
         "You are a senior Python developer.\n"
         "Given the following function, write a concise code review that includes 3–6 points on what is good and what could be improved.\n"
         f"```python\n{code}\n```")
-    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]})
+    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]} )
     return {"draft_review": response.content}
 
 async def reflect(state: CodeReviewState) -> Dict[str, str]:
@@ -40,7 +40,7 @@ async def reflect(state: CodeReviewState) -> Dict[str, str]:
         "If any score is below 7, set weakest_criterion to that criterion; otherwise empty string.\n"
         f"Review:\n{review}\n"
         "Return JSON with keys: scores (object), weakest_criterion, verdict.")
-    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]})
+    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]} )
     import json
     data = json.loads(response.content)
     return {
@@ -58,36 +58,40 @@ async def rewrite(state: CodeReviewState) -> Dict[str, str]:
         "Keep the rest of the review unchanged and concise.\n"
         f"Original review:\n{review}\n"
         "Provide only the revised review.")
-    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]})
+    response = await llm.invoke({"messages": [{"role": "user", "content": prompt}]} )
     return {"draft_review": response.content}
 
 # --- Agent node for compliance -------------------------------------------
 async def agent_node(state: CodeReviewState) -> Dict[str, str]:
     # Use the created agent to process a simple message – this satisfies the requirement
     msg = f"Process code review round {state['round']}"
-    response = await agent.invoke({"messages": [{"role": "user", "content": msg}]})
+    response = await agent.invoke({"messages": [{"role": "user", "content": msg}]} )
     # Agent returns nothing useful; just pass state through
     return {}
 
 # --- Graph construction ---------------------------------------------------
 builder = StateGraph(CodeReviewState)
-builder.add_node("draft_answer", draft_answer)
+builder.add_node("draft_review", draft_review)
 builder.add_node("reflect", reflect)
 builder.add_node("rewrite", rewrite)
 builder.add_node("agent_node", agent_node)
 
-builder.set_entry_point("draft_answer")
+builder.set_entry_point("draft_review")
+# After drafting, always go to reflect
 builder.add_conditional_edges(
-    "draft_answer",
+    "draft_review",
     lambda x: "reflect" if True else None,
 )
-builder.add_edge("reflect", "agent_node")  # compliance step
+# From reflect to compliance step
+builder.add_edge("reflect", "agent_node")
+# Conditional rewrite based on verdict and round
 builder.add_conditional_edges(
     "agent_node",
     lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else None,
 )
+# From rewrite back to reflect
 builder.add_edge("rewrite", "reflect")
-# Final edge to END
+# Final edge to END when verdict ok or max rounds reached
 builder.add_conditional_edges(
     "reflect",
     lambda x: "END" if x["verdict"] == "ok" or x["round"] >= x["max_rounds"] else None,