Solution published successfully: update main.py

main.py

@@ -2,8 +2,10 @@
 
 Implementation follows assignment:
 - State: CodeReviewState with 4 criteria.
-- Nodes: review_and_critique, rewrite.
-- Graph: START -> review_and_critique -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> review_and_critique).
+- Nodes: draft_review, reflect, rewrite.
+- Graph: START -> draft_review -> reflect
+  - ok -> END
+  - needs_revision & round < max_rounds -> rewrite -> reflect
 - Uses LangGraph and LangChain OpenAI for LLM calls.
 - Structured output for critique via Pydantic model.
 - Demo function sort_numbers.
@@ -11,7 +13,6 @@ Implementation follows assignment:
 
 from __future__ import annotations
 
-import os
 from typing import TypedDict, Dict
 
 from langgraph.graph import StateGraph, END
@@ -30,18 +31,19 @@ class CodeReviewState(TypedDict):
     max_rounds: int
 
 # ---------- LLM ----------
-# Use OpenAI only
 llm = ChatOpenAI(temperature=0)
 
 # ---------- Nodes ----------
-class CritiqueOutput(BaseModel):
+class DraftReviewOutput(BaseModel):
     review: str = Field(..., description="Draft review text")
+
+class ReflectOutput(BaseModel):
     scores: Dict[str, int] = Field(..., description="Scores 0-10 for each criterion")
     weakest_criterion: str = Field(..., description="Criterion with lowest score")
     verdict: str = Field(..., description="'ok' or 'needs_revision'")
 
 
-def review_and_critique(state: CodeReviewState) -> CodeReviewState:
+def draft_review(state: CodeReviewState) -> CodeReviewState:
     code = state["code"]
     prompt = f"""
     Write a concise code review (3-6 bullet points) for the following Python function. Focus on style, correctness, and potential improvements.
@@ -49,13 +51,23 @@ def review_and_critique(state: CodeReviewState) -> CodeReviewState:
     ```python
     {code}
     ```
-
-    Then evaluate the review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10.
-    Return a JSON object with keys: review (string), scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
     """
     response = llm.invoke([HumanMessage(content=prompt)])
-    data = CritiqueOutput.model_validate_json(response.content)
+    data = DraftReviewOutput.model_validate_json(response.content)
     state["draft_review"] = data.review
+    return state
+
+
+def reflect(state: CodeReviewState) -> CodeReviewState:
+    review = state["draft_review"]
+    prompt = f"""
+    Evaluate the following review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10. Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
+
+    Review:
+    {review}
+    """
+    response = llm.invoke([HumanMessage(content=prompt)])
+    data = ReflectOutput.model_validate_json(response.content)
     state["criteria_scores"] = data.scores
     state["weakest_criterion"] = data.weakest_criterion
     state["verdict"] = data.verdict
@@ -77,22 +89,26 @@ def rewrite(state: CodeReviewState) -> CodeReviewState:
 
 # ---------- Graph ----------
 builder = StateGraph(CodeReviewState)
-builder.add_node("review_and_critique", review_and_critique)
+builder.add_node("draft_review", draft_review)
+builder.add_node("reflect", reflect)
 builder.add_node("rewrite", rewrite)
 
-builder.set_entry_point("review_and_critique")
-# After initial review_and_critique, decide to end if verdict ok
+builder.set_entry_point("draft_review")
+# After draft_review, go to reflect
+builder.add_edge("draft_review", "reflect")
+# After reflect, decide
 builder.add_conditional_edges(
-    "review_and_critique",
+    "reflect",
     lambda state: state["verdict"] == "ok",
     {"ok": END, "needs_revision": "rewrite"},
 )
-# After rewrite, go back to review_and_critique if rounds remain
+# After rewrite, go back to reflect if rounds remain
 builder.add_conditional_edges(
     "rewrite",
-    lambda state: "review_and_critique" if state["round"] < state["max_rounds"] else END,
-    {"review_and_critique": "review_and_critique", END: END}
+    lambda state: state["round"] < state["max_rounds"],
+    {"continue": "reflect", "end": END},
 )
+# Compile graph
 
 graph = builder.compile()