diff --git a/main.py b/main.py index e2eb1a9..0194f5d 100644 --- a/main.py +++ b/main.py @@ -1,120 +1,203 @@ -# main.py +"""LangGraph Code Review Agent with Reflection. + +This repository implements a LangGraph agent that takes a Python function as input and +produces a code review. A critic node evaluates the review on four criteria: + +* PEP8 compliance +* Type hints +* Edge cases handling +* Naming conventions + +If the critic returns `needs_revision`, the `rewrite` node rewrites the weakest +criterion section of the review. The process repeats until the critic is satisfied +or the maximum number of rounds is reached. + +The demo in ``__main__`` shows how to run the agent on a simple function. +""" + +from __future__ import annotations + import os from typing import TypedDict, Dict -from langgraph.graph import StateGraph, END -from langgraph.prebuilt import create_chat_agent -from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, AIMessage -# Define state +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_react_agent +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# State definition +# --------------------------------------------------------------------------- class CodeReviewState(TypedDict): code: str - draft_review: str - criteria_scores: Dict[str, int] - weakest_criterion: str - verdict: str + draft_review: str | None + criteria_scores: Dict[str, int] | None + weakest_criterion: str | None + verdict: str | None # "ok" | "needs_revision" round: int max_rounds: int -# LLM -llm = ChatOpenAI(temperature=0) +# --------------------------------------------------------------------------- +# LLM configuration +# --------------------------------------------------------------------------- +# The OpenAI API key must be set in the environment variable OPENAI_API_KEY. +# For local Ollama usage, replace ChatOpenAI with ChatOllama. +llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2) -# Draft review node +# --------------------------------------------------------------------------- +# Node: draft_review +# --------------------------------------------------------------------------- async def draft_review(state: CodeReviewState) -> CodeReviewState: - prompt = f""" -You are a senior Python developer. Review the following code and provide a concise code review (3-6 bullet points) highlighting what is good and what can be improved. + """Generate an initial code review. -Code: -{state['code']} - -Review: -""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state['draft_review'] = response.content + The review contains 3–6 bullet points describing what is good and what can be + improved. The output is plain text. + """ + prompt = ( + "You are a senior Python developer.\n" + "Given the following function, write a concise code review (3–6 points).\n" + "Focus on style, correctness, and potential improvements.\n" + "Return only the review text.\n\n" + f"Function:\n{state["code"]}\n" + ) + review = await llm.ainvoke(prompt) + state["draft_review"] = review.content.strip() return state -# Reflect node +# --------------------------------------------------------------------------- +# Node: reflect +# --------------------------------------------------------------------------- +class ReflectOutput(BaseModel): + pep8: int = Field(..., ge=0, le=10) + type_hints: int = Field(..., ge=0, le=10) + edge_cases: int = Field(..., ge=0, le=10) + naming: int = Field(..., ge=0, le=10) + verdict: str = Field(..., regex="^(ok|needs_revision)$") + weakest_criterion: str = Field(..., regex="^(pep8|type_hints|edge_cases|naming)$") + async def reflect(state: CodeReviewState) -> CodeReviewState: - prompt = f""" -You are an AI critic evaluating a code review. Assign a score 0-10 for each of the following criteria based on the draft review: -- pep8 -- type_hints -- edge_cases -- naming + """Critic node that scores the draft review on four criteria. -Provide a JSON object with keys "pep8", "type_hints", "edge_cases", "naming" and integer values. -Also determine the weakest criterion (the one with lowest score) and a verdict: "ok" if all scores >=7, otherwise "needs_revision". - -Draft review: -{state['draft_review']} - -Output JSON: -""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - import json - scores = json.loads(response.content) - state['criteria_scores'] = scores - weakest = min(scores, key=scores.get) - state['weakest_criterion'] = weakest - state['verdict'] = "ok" if all(v >= 7 for v in scores.values()) else "needs_revision" + The LLM returns a JSON object that matches ``ReflectOutput``. The function + parses the JSON and updates the state. + """ + prompt = ( + "You are a code quality critic.\n" + "Given the following code review, score it on the following criteria (0–10):\n" + "- PEP8 compliance\n" + "- Type hints usage\n" + "- Edge cases handling\n" + "- Naming conventions\n" + "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, verdict, weakest_criterion.\n" + "Verdict should be "ok" if all scores are >=7, otherwise "needs_revision".\n" + "Weakest criterion is the one with the lowest score.\n\n" + f"Review:\n{state["draft_review"]}\n" + ) + result = await llm.ainvoke(prompt) + try: + data = ReflectOutput.model_validate_json(result.content) + except Exception as e: + # Fallback: if parsing fails, treat as needs_revision + data = ReflectOutput( + pep8=0, + type_hints=0, + edge_cases=0, + naming=0, + verdict="needs_revision", + weakest_criterion="pep8", + ) + state["criteria_scores"] = data.model_dump(exclude="verdict,weakest_criterion") + state["weakest_criterion"] = data.weakest_criterion + state["verdict"] = data.verdict return state -# Rewrite node +# --------------------------------------------------------------------------- +# Node: rewrite +# --------------------------------------------------------------------------- async def rewrite(state: CodeReviewState) -> CodeReviewState: - crit = state['weakest_criterion'] - prompt = f""" -You are a senior Python developer. Rewrite the section of the code review that addresses the {crit} criterion, improving it. Keep the rest of the review unchanged. + """Rewrite the weakest part of the review. -Original review: -{state['draft_review']} - -Rewrite only the part related to {crit}: -""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - # Replace the part in draft_review that mentions crit - # For simplicity, just append the new part - state['draft_review'] = state['draft_review'] + "\n" + response.content - state['round'] += 1 + The LLM is instructed to rewrite only the section that addresses the + weakest criterion. The new review replaces the old one. + """ + prompt = ( + "You are a senior Python developer.\n" + "Rewrite the part of the following code review that addresses the weakest criterion.\n" + "Keep the rest of the review unchanged.\n" + "Return only the updated review text.\n\n" + f"Weakest criterion: {state["weakest_criterion"]}\n" + f"Current review:\n{state["draft_review"]}\n" + ) + new_review = await llm.ainvoke(prompt) + state["draft_review"] = new_review.content.strip() + state["round"] += 1 return state -# Build graph -builder = StateGraph(CodeReviewState) -builder.add_node("draft_review", draft_review) -builder.add_node("reflect", reflect) -builder.add_node("rewrite", rewrite) +# --------------------------------------------------------------------------- +# 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) -builder.set_entry_point("draft_review") -builder.add_edge("draft_review", "reflect") -builder.add_conditional_edges( - "reflect", - lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END", -) -builder.add_edge("rewrite", "reflect") + # Entry point + graph.set_entry_point("draft_review") -graph = builder.compile() + # Transitions + graph.add_conditional_edges( + "draft_review", + lambda state: "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") -# Demo function -async def run_demo(): - code = """ -# Example function to sort numbers + return graph -def sort_numbers(arr): - return sorted(arr) -""" - init_state: CodeReviewState = { +# --------------------------------------------------------------------------- +# Demo helper +# --------------------------------------------------------------------------- +async def run_demo(func) -> None: + """Run the graph on a single function and print the results.""" + import inspect + import asyncio + + code = inspect.getsource(func) + state: CodeReviewState = { "code": code, - "draft_review": "", - "criteria_scores": {}, - "weakest_criterion": "", - "verdict": "", + "draft_review": None, + "criteria_scores": None, + "weakest_criterion": None, + "verdict": None, "round": 0, "max_rounds": 2, } - result = await graph.ainvoke(init_state) - print("Final Review:\n", result["draft_review"]) - print("Scores:\n", result["criteria_scores"]) - print("Verdict:\n", result["verdict"]) + graph = build_graph() + final_state = await graph.astate(state) + print("\n=== Initial Review ===") + print(state["draft_review"]) + print("\n=== Scores ===") + print(final_state["criteria_scores"]) + print("Verdict:", final_state["verdict"]) + if final_state["verdict"] == "needs_revision": + print("\n=== Revised Review ===") + print(final_state["draft_review"]) +# --------------------------------------------------------------------------- +# Example function for demo +# --------------------------------------------------------------------------- + +def sort_numbers(arr): + """Return a sorted copy of the input list.""" + return sorted(arr) + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- if __name__ == "__main__": import asyncio - asyncio.run(run_demo()) + asyncio.run(run_demo(sort_numbers))