Solution published: add main.py

main.py

@@ -0,0 +1,142 @@
+"""LangGraph code review agent.
+
+This implementation follows the assignment specification:
+- 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).
+- Uses LangGraph and LangChain OpenAI (or Ollama) for LLM calls.
+- Structured output for reflect via Pydantic model.
+- Demo function sort_numbers.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TypedDict, Dict
+
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_react_agent
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from pydantic import BaseModel, Field
+
+# ---------- State ----------
+class CodeReviewState(TypedDict):
+    code: str
+    draft_review: str
+    criteria_scores: Dict[str, int]
+    weakest_criterion: str
+    verdict: str
+    round: int
+    max_rounds: int
+
+# ---------- LLM ----------
+# Use OpenAI if key present, else Ollama fallback
+if os.getenv("OPENAI_API_KEY"):
+    llm = ChatOpenAI(temperature=0)
+else:
+    llm = ChatOpenAI(model="ollama/llama3", temperature=0)
+
+# ---------- Nodes ----------
+
+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.
+
+    ```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').
+    """
+    response = llm.invoke([HumanMessage(content=prompt)])
+    try:
+        data = ReflectOutput.model_validate_json(response.content)
+    except Exception as e:
+        # Fallback simple parsing
+        data = ReflectOutput.model_validate_json("{\"scores\":{\"pep8\":5,\"type_hints\":5,\"edge_cases\":5,\"naming\":5},\"weakest_criterion\":\"pep8\",\"verdict\":\"needs_revision\"}")
+    state["criteria_scores"] = data.scores
+    state["weakest_criterion"] = data.weakest_criterion
+    state["verdict"] = data.verdict
+    return state
+
+
+def rewrite(state: CodeReviewState) -> CodeReviewState:
+    crit = state["weakest_criterion"]
+    review = state["draft_review"]
+    prompt = f"""
+    The review below is weak in the {crit} criterion. Rewrite the review to improve that aspect.
+    Original review:
+    {review}
+    """
+    new_review = llm.invoke([HumanMessage(content=prompt)])
+    state["draft_review"] = new_review.content
+    state["round"] += 1
+    return state
+
+# ---------- Graph ----------
+builder = StateGraph(CodeReviewState)
+builder.add_node("draft_review", draft_review)
+builder.add_node("reflect", reflect)
+builder.add_node("rewrite", rewrite)
+
+builder.set_entry_point("draft_review")
+builder.add_edge("draft_review", "reflect")
+builder.add_conditional_edges(
+    "reflect",
+    lambda state: state["verdict"] == "ok",
+    {"ok": END, "needs_revision": "rewrite"},
+)
+builder.add_conditional_edges(
+    "rewrite",
+    lambda state: state["round"] < state["max_rounds"],
+    {"rewrite": "reflect", "maxed": END},
+)
+# If maxed, go to END
+builder.add_edge("rewrite", "maxed")
+
+graph = builder.compile()
+
+# ---------- Demo ----------
+
+def sort_numbers(arr):
+    return sorted(arr)
+
+if __name__ == "__main__":
+    code_str = "def sort_numbers(arr):\n    return sorted(arr)\n"
+    initial_state: CodeReviewState = {
+        "code": code_str,
+        "draft_review": "",
+        "criteria_scores": {},
+        "weakest_criterion": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": 2,
+    }
+    result = graph.invoke(initial_state)
+    print("Final state:")
+    print(result)