Initial implementation of LangGraph code review agent with reflection and rewrite loop.: add repo/main.py

repo/main.py

@@ -0,0 +1,115 @@
+"""
+LangGraph Code Review Agent
+Author: Auto-generated
+"""
+import os
+from typing import TypedDict, Dict
+from langgraph.graph import StateGraph
+# from langgraph.prebuilt import create_agent_executor  # unused
+from langchain_openai import ChatOpenAI
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# ---------- State ----------
+class CodeReviewState(TypedDict):
+    code: str
+    draft_review: str
+    criteria_scores: Dict[str, int]
+    weakest_criterion: str
+    verdict: str  # "ok" | "needs_revision"
+    round: int
+    max_rounds: int
+
+# ---------- Nodes ----------
+async def draft_review(state: CodeReviewState) -> CodeReviewState:
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
+    prompt = (
+        "You are a senior Python developer.\n"
+        f"Given the following code:\n{state['code']}\n"
+        "Write a concise code review (3–6 points). Do not include any other text."  # noqa: E501
+    )
+    response = await llm.invoke(prompt)
+    state["draft_review"] = response.content.strip()
+    return state
+
+async def reflect(state: CodeReviewState) -> CodeReviewState:
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)
+    prompt = f'''You are an automated code review critic.
+Draft review:
+{state['draft_review']}
+Score the draft on the following criteria (0–10):
+- PEP8 compliance
+- Type hints usage
+- Edge case handling
+- Naming conventions
+Return a JSON object with keys: pep8, type_hints, edge_cases, naming. Also provide the weakest criterion and verdict ("ok" if all scores >=7 else "needs_revision").'''
+    response = await llm.invoke(prompt)
+    import json
+    try:
+        data = json.loads(response.content.strip())
+    except Exception as e:
+        raise ValueError(f"Failed to parse critic output: {e}\n{response.content}")
+    state["criteria_scores"] = {
+        "pep8": int(data.get("pep8", 0)),
+        "type_hints": int(data.get("type_hints", 0)),
+        "edge_cases": int(data.get("edge_cases", 0)),
+        "naming": int(data.get("naming", 0)),
+    }
+    state["weakest_criterion"] = data.get("weakest_criterion", "")
+    state["verdict"] = data.get("verdict", "needs_revision")
+    return state
+
+async def rewrite(state: CodeReviewState) -> CodeReviewState:
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
+    prompt = (
+        f"The draft review was insufficient in the {state['weakest_criterion']} area.\n"
+        "Rewrite only that part of the review to improve it, keeping other points unchanged.\n"
+        f"Current draft:\n{state['draft_review']}\n"
+        "Provide the updated draft."  # noqa: E501
+    )
+    response = await llm.invoke(prompt)
+    state["draft_review"] = response.content.strip()
+    state["round"] += 1
+    return state
+
+# ---------- Graph ----------
+builder = StateGraph(CodeReviewState)
+builder.add_node("draft", draft_review)
+builder.add_node("reflect", reflect)
+builder.add_node("rewrite", rewrite)
+
+builder.set_entry_point("draft")
+builder.add_edge("draft", "reflect")
+builder.add_conditional_edges(
+    "reflect",
+    lambda state: (
+        "end"
+        if state["verdict"] == "ok"
+        else ("rewrite" if state["round"] < state.get("max_rounds", 2) else "end")
+    ),
+)
+builder.add_edge("rewrite", "reflect")
+
+graph = builder.compile()
+
+# ---------- Demo ----------
+if __name__ == "__main__":
+    sample_code = (
+        "def sort_numbers(arr):\n"
+        "    return sorted(arr)"
+    )
+    initial_state: CodeReviewState = {
+        "code": sample_code,
+        "draft_review": "",
+        "criteria_scores": {},
+        "weakest_criterion": "",
+        "verdict": "needs_revision",
+        "round": 0,
+        "max_rounds": 2,
+    }
+    result = graph.invoke(initial_state)
+    print("\n--- Final Review ---")
+    print(result["draft_review"])
+    print("\nScores:", result["criteria_scores"])
+    print("Verdict:", result["verdict"])