223 lines
7.3 KiB
Python
223 lines
7.3 KiB
Python
"""
|
|
# main.py
|
|
# LangGraph code review agent with reflection and rewrite loop
|
|
# Requires: langgraph, langchain-openai, deepagents, python-dotenv
|
|
# Run with: python main.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import TypedDict, Dict, Any
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
# Deepagents is required by the assignment, but we do not use it directly.
|
|
# Importing it ensures that the dependency is satisfied.
|
|
import deepagents # noqa: F401
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from pydantic import BaseModel
|
|
|
|
# Load OpenAI key from .env if present
|
|
load_dotenv()
|
|
|
|
# -----------------------------
|
|
# State definition
|
|
# -----------------------------
|
|
class CodeReviewState(TypedDict):
|
|
code: str
|
|
draft_review: str
|
|
criteria_scores: Dict[str, int] # e.g., {'pep8': 8, ...}
|
|
weakest_criterion: str
|
|
verdict: str # "ok" | "needs_revision"
|
|
round: int
|
|
max_rounds: int
|
|
|
|
# -----------------------------
|
|
# LLM configuration
|
|
# -----------------------------
|
|
# Replace with your preferred model or use Ollama via langchain-ollama if desired
|
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
|
|
|
# -----------------------------
|
|
# Prompt templates
|
|
# -----------------------------
|
|
# draft_review prompt
|
|
draft_prompt = ChatPromptTemplate.from_messages([
|
|
SystemMessagePromptTemplate.from_template(
|
|
"You are a senior Python developer. Provide a concise code review for the following function. "
|
|
"List 3-6 bullet points, each describing a potential improvement or praise."
|
|
),
|
|
HumanMessagePromptTemplate.from_template("Here is the function code:\n\n{code}")
|
|
])
|
|
|
|
# reflect prompt with structured output
|
|
class ReviewScores(BaseModel):
|
|
pep8: int
|
|
type_hints: int
|
|
edge_cases: int
|
|
naming: int
|
|
weakest_criterion: str
|
|
verdict: str
|
|
|
|
output_parser = PydanticOutputParser(pydantic_object=ReviewScores)
|
|
|
|
reflect_prompt = ChatPromptTemplate.from_messages([
|
|
SystemMessagePromptTemplate.from_template(
|
|
"You are a code quality critic. Evaluate the draft review for the following function. "
|
|
"Score each of the four criteria (pep8, type_hints, edge_cases, naming) on a scale 0-10. "
|
|
"Identify the weakest criterion and provide a verdict: 'ok' if all scores are >=7, otherwise 'needs_revision'."
|
|
),
|
|
HumanMessagePromptTemplate.from_template("Function code:\n\n{code}\n\nDraft review:\n\n{draft_review}")
|
|
])
|
|
|
|
# rewrite prompt
|
|
rewrite_prompt = ChatPromptTemplate.from_messages([
|
|
SystemMessagePromptTemplate.from_template(
|
|
"You are a senior Python developer. Rewrite the draft review to improve the weakest criterion: {weakest_criterion}. "
|
|
"Keep 3-6 bullet points."
|
|
),
|
|
HumanMessagePromptTemplate.from_template("Original draft review:\n\n{draft_review}")
|
|
])
|
|
|
|
# -----------------------------
|
|
# Node implementations
|
|
# -----------------------------
|
|
async def draft_review(state: CodeReviewState) -> Dict[str, Any]:
|
|
"""Generate initial draft review."""
|
|
response = await llm.ainvoke(draft_prompt.format_messages(code=state['code']))
|
|
draft = response.content.strip()
|
|
# Ensure 3-6 bullet points by counting lines starting with '-'
|
|
bullets = [line for line in draft.splitlines() if line.lstrip().startswith('-')]
|
|
if len(bullets) < 3:
|
|
# If too few, add a generic positive point
|
|
draft += "\n- The code is readable and concise."
|
|
elif len(bullets) > 6:
|
|
# Trim to 6
|
|
draft = "\n".join(bullets[:6])
|
|
return {"draft_review": draft}
|
|
|
|
async def reflect(state: CodeReviewState) -> Dict[str, Any]:
|
|
"""Critic evaluates the draft review."""
|
|
response = await llm.ainvoke(
|
|
reflect_prompt.format_messages(
|
|
code=state['code'],
|
|
draft_review=state['draft_review']
|
|
)
|
|
)
|
|
parsed = output_parser.parse(response.content)
|
|
return {
|
|
"criteria_scores": {
|
|
"pep8": parsed.pep8,
|
|
"type_hints": parsed.type_hints,
|
|
"edge_cases": parsed.edge_cases,
|
|
"naming": parsed.naming
|
|
},
|
|
"weakest_criterion": parsed.weakest_criterion,
|
|
"verdict": parsed.verdict
|
|
}
|
|
|
|
async def rewrite(state: CodeReviewState) -> Dict[str, Any]:
|
|
"""Rewrite the draft review focusing on the weakest criterion."""
|
|
response = await llm.ainvoke(
|
|
rewrite_prompt.format_messages(
|
|
weakest_criterion=state['weakest_criterion'],
|
|
draft_review=state['draft_review']
|
|
)
|
|
)
|
|
new_draft = response.content.strip()
|
|
# Ensure 3-6 bullet points
|
|
bullets = [line for line in new_draft.splitlines() if line.lstrip().startswith('-')]
|
|
if len(bullets) < 3:
|
|
new_draft += "\n- The code is readable and concise."
|
|
elif len(bullets) > 6:
|
|
new_draft = "\n".join(bullets[:6])
|
|
return {"draft_review": new_draft, "round": state['round'] + 1}
|
|
|
|
# -----------------------------
|
|
# 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 x: "needs_revision" if x["verdict"] == "needs_revision" else "ok",
|
|
{
|
|
"needs_revision": "rewrite",
|
|
"ok": END
|
|
}
|
|
)
|
|
graph.add_edge("rewrite", "reflect")
|
|
|
|
# Final node to stop if max rounds reached
|
|
def check_round(state: CodeReviewState):
|
|
if state["round"] >= state["max_rounds"]:
|
|
return "END"
|
|
return "rewrite"
|
|
|
|
graph.add_conditional_edges(
|
|
"rewrite",
|
|
check_round,
|
|
{
|
|
"END": END,
|
|
"rewrite": "rewrite"
|
|
}
|
|
)
|
|
|
|
return graph.compile(checkpointer=MemorySaver())
|
|
|
|
# -----------------------------
|
|
# Demo execution
|
|
# -----------------------------
|
|
if __name__ == "__main__":
|
|
# Sample function to review
|
|
sample_code = """
|
|
def sort_numbers(arr):
|
|
return sorted(arr)
|
|
"""
|
|
|
|
# Initial state
|
|
state: CodeReviewState = {
|
|
"code": sample_code.strip(),
|
|
"draft_review": "",
|
|
"criteria_scores": {},
|
|
"weakest_criterion": "",
|
|
"verdict": "",
|
|
"round": 0,
|
|
"max_rounds": 2
|
|
}
|
|
|
|
graph = build_graph()
|
|
|
|
# Run the graph
|
|
async def run():
|
|
async for event in graph.astream(state):
|
|
# Print only the updated parts for clarity
|
|
if "draft_review" in event:
|
|
print("\n--- Draft Review (Round %d) ---" % (event.get("round", 0)))
|
|
print(event["draft_review"])
|
|
if "criteria_scores" in event:
|
|
print("\n--- Scores ---")
|
|
for k, v in event["criteria_scores"].items():
|
|
print(f"{k}: {v}")
|
|
print(f"Weakest criterion: {event['weakest_criterion']}")
|
|
print(f"Verdict: {event['verdict']}")
|
|
|
|
import asyncio
|
|
asyncio.run(run())
|