116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
"""
|
||
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"])
|