diff --git a/main.py b/main.py new file mode 100644 index 0000000..12cba81 --- /dev/null +++ b/main.py @@ -0,0 +1,198 @@ +""" +LangGraph Code Review Agent +=========================== + +This repository contains a minimal LangGraph implementation that +performs a code review on a Python function. The graph consists of +three nodes: + +* ``draft_review`` – generates an initial review. +* ``reflect`` – a critic that scores the review on four criteria + (PEP8, type hints, edge cases, naming) and decides whether a + rewrite is required. +* ``rewrite`` – rewrites the weakest part of the review. + +The graph runs for a maximum of ``max_rounds`` (default 2). The +demo can be executed with ``python -m main``. +""" + +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 + +# --------------------------------------------------------------------------- +# State definition +# --------------------------------------------------------------------------- + +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 + +# --------------------------------------------------------------------------- +# LLM configuration +# --------------------------------------------------------------------------- + +# The user must set the OPENAI_API_KEY environment variable. +llm = ChatOpenAI(temperature=0.0, model="gpt-4o-mini") + +# --------------------------------------------------------------------------- +# Node implementations +# --------------------------------------------------------------------------- + +async def draft_review(state: CodeReviewState) -> CodeReviewState: + """Generate an initial review of the provided code. + + The review is a short list of 3–6 bullet points. + """ + code = state["code"] + prompt = ( + "You are a senior Python developer.\n" + "Review the following function and provide a concise list of 3–6 points\n" + "highlighting what is good and what could be improved.\n" + "Do not mention the criteria – just give the review.\n" + f"Function:\n{code}\n" + "Review:" # LLM will continue after this + ) + review = await llm.ainvoke([HumanMessage(content=prompt)]) + state["draft_review"] = review.content.strip() + return state + +async def reflect(state: CodeReviewState) -> CodeReviewState: + """Critic that scores the draft review on four criteria. + + The output is a JSON object with keys: + * pep8, type_hints, edge_cases, naming – integers 0‑10 + * weakest_criterion – one of the four keys + * verdict – "ok" or "needs_revision" + """ + review = state["draft_review"] + code = state["code"] + prompt = ( + "You are a code quality critic.\n" + "Given the following code and its draft review, score the review on\n" + "four criteria: PEP8, type hints, edge cases, naming.\n" + "Return a JSON object with keys: pep8, type_hints, edge_cases, naming,\n" + "weakest_criterion, verdict.\n" + f"Code:\n{code}\n" + f"Draft review:\n{review}\n" + "Answer in JSON only." + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + import json + try: + scores = json.loads(response.content) + except Exception as e: + # Fallback: if parsing fails, treat as needs_revision + scores = { + "pep8": 0, + "type_hints": 0, + "edge_cases": 0, + "naming": 0, + "weakest_criterion": "pep8", + "verdict": "needs_revision", + } + state["criteria_scores"] = { + "pep8": int(scores.get("pep8", 0)), + "type_hints": int(scores.get("type_hints", 0)), + "edge_cases": int(scores.get("edge_cases", 0)), + "naming": int(scores.get("naming", 0)), + } + state["weakest_criterion"] = scores.get("weakest_criterion", "pep8") + state["verdict"] = scores.get("verdict", "needs_revision") + return state + +async def rewrite(state: CodeReviewState) -> CodeReviewState: + """Rewrite the weakest part of the review. + + The node receives the current state and the weakest criterion. + It generates a new review that specifically addresses that criterion. + """ + code = state["code"] + weakest = state["weakest_criterion"] + prompt = ( + "You are a senior Python developer.\n" + "Rewrite the draft review to improve the part related to the following criterion: " + f"{weakest}.\n" + "Keep the rest of the review unchanged.\n" + "Output only the updated review.\n" + f"Current draft review:\n{state['draft_review']}\n" + ) + new_review = await llm.ainvoke([HumanMessage(content=prompt)]) + state["draft_review"] = new_review.content.strip() + state["round"] += 1 + return state + +# --------------------------------------------------------------------------- +# Graph construction +# --------------------------------------------------------------------------- + +def build_graph() -> StateGraph[CodeReviewState]: + graph = StateGraph(CodeReviewState) + graph.add_node("draft_review", draft_review) + graph.add_node("reflect", reflect) + graph.add_node("rewrite", rewrite) + + # Entry point + graph.set_entry_point("draft_review") + + # Transitions + graph.add_edge("draft_review", "reflect") + graph.add_conditional_edges( + "reflect", + lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END", + ) + graph.add_edge("rewrite", "reflect") + + return graph + +# --------------------------------------------------------------------------- +# Demo execution +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import argparse + import textwrap + + parser = argparse.ArgumentParser(description="Run the code review graph on a demo function.") + parser.add_argument("--max-rounds", type=int, default=2, help="Maximum number of rewrite rounds") + args = parser.parse_args() + + # Demo function – can be replaced by any user code + demo_code = textwrap.dedent( + """ + def sort_numbers(arr): + return sorted(arr) + """ + ).strip() + + initial_state: CodeReviewState = { + "code": demo_code, + "draft_review": "", + "criteria_scores": {}, + "weakest_criterion": "", + "verdict": "", + "round": 0, + "max_rounds": args.max_rounds, + } + + graph = build_graph() + final_state = graph.invoke(initial_state) + + print("\n=== Final Review ===") + print(final_state["draft_review"]) + print("\n=== Scores ===") + for k, v in final_state["criteria_scores"].items(): + print(f"{k}: {v}") + print(f"Verdict: {final_state['verdict']}") + print(f"Rounds: {final_state['round']}")