Solution ready for publishing: update main.py

main.py

@@ -2,10 +2,10 @@
 
 Implementation follows assignment:
 - State: CodeReviewState with 4 criteria.
-- Nodes: draft_review, reflect, rewrite.
-- Graph: START -> draft_review -> reflect -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> reflect).
+- Nodes: review_and_critique, rewrite.
+- Graph: START -> review_and_critique -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> review_and_critique).
 - Uses LangGraph and LangChain OpenAI for LLM calls.
-- Structured output for reflect via Pydantic model.
+- Structured output for critique via Pydantic model.
 - Demo function sort_numbers.
 """
 
@@ -34,8 +34,14 @@ class CodeReviewState(TypedDict):
 llm = ChatOpenAI(temperature=0)
 
 # ---------- Nodes ----------
+class CritiqueOutput(BaseModel):
+    review: str = Field(..., description="Draft review text")
+    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 draft_review(state: CodeReviewState) -> CodeReviewState:
+
+def review_and_critique(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.
@@ -43,34 +49,13 @@ def draft_review(state: CodeReviewState) -> CodeReviewState:
     ```python
     {code}
     ```
-    """
-    review = llm.invoke([HumanMessage(content=prompt)])
-    state["draft_review"] = review.content
-    return state
 
-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 reflect(state: CodeReviewState) -> CodeReviewState:
-    code = state["code"]
-    review = state["draft_review"]
-    prompt = f"""
-    You are a code review critic. Evaluate the following review of a Python function.
-    Function code:
-    ```python
-    {code}
-    ```
-    Review:
-    {review}
-
-    Score the 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').
+    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 = ReflectOutput.model_validate_json(response.content)
+    data = CritiqueOutput.model_validate_json(response.content)
+    state["draft_review"] = data.review
     state["criteria_scores"] = data.scores
     state["weakest_criterion"] = data.weakest_criterion
     state["verdict"] = data.verdict
@@ -92,22 +77,21 @@ def rewrite(state: CodeReviewState) -> CodeReviewState:
 
 # ---------- Graph ----------
 builder = StateGraph(CodeReviewState)
-builder.add_node("draft_review", draft_review)
-builder.add_node("reflect", reflect)
+builder.add_node("review_and_critique", review_and_critique)
 builder.add_node("rewrite", rewrite)
 
-builder.set_entry_point("draft_review")
-builder.add_edge("draft_review", "reflect")
+builder.set_entry_point("review_and_critique")
+# After initial review_and_critique, decide to end if verdict ok
 builder.add_conditional_edges(
-    "reflect",
+    "review_and_critique",
     lambda state: state["verdict"] == "ok",
     {"ok": END, "needs_revision": "rewrite"},
 )
-# After rewrite, decide to reflect again or end if max rounds reached
+# After rewrite, go back to review_and_critique if rounds remain
 builder.add_conditional_edges(
     "rewrite",
-    lambda state: "reflect" if state["round"] < state["max_rounds"] else END,
-    {"reflect": "reflect", END: END}
+    lambda state: "review_and_critique" if state["round"] < state["max_rounds"] else END,
+    {"review_and_critique": "review_and_critique", END: END}
 )
 
 graph = builder.compile()