199 lines
6.8 KiB
Python
199 lines
6.8 KiB
Python
"""
|
||
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']}")
|