204 lines
7.5 KiB
Python
204 lines
7.5 KiB
Python
"""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 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 | None
|
||
criteria_scores: Dict[str, int] | None
|
||
weakest_criterion: str | None
|
||
verdict: str | None # "ok" | "needs_revision"
|
||
round: int
|
||
max_rounds: int
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: draft_review
|
||
# ---------------------------------------------------------------------------
|
||
async def draft_review(state: CodeReviewState) -> CodeReviewState:
|
||
"""Generate an initial code review.
|
||
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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:
|
||
"""Critic node that scores the draft review on four criteria.
|
||
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: rewrite
|
||
# ---------------------------------------------------------------------------
|
||
async def rewrite(state: CodeReviewState) -> CodeReviewState:
|
||
"""Rewrite the weakest part of the review.
|
||
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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_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")
|
||
|
||
return graph
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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": None,
|
||
"criteria_scores": None,
|
||
"weakest_criterion": None,
|
||
"verdict": None,
|
||
"round": 0,
|
||
"max_rounds": 2,
|
||
}
|
||
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(sort_numbers))
|