134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
"""LangGraph code review agent.
|
|
|
|
Implementation follows assignment:
|
|
- State: CodeReviewState with 4 criteria.
|
|
- Nodes: draft_review, reflect, rewrite.
|
|
- Graph: START -> draft_review -> reflect
|
|
- ok -> END
|
|
- needs_revision & round < max_rounds -> rewrite -> reflect
|
|
- Uses LangGraph and LangChain OpenAI for LLM calls.
|
|
- Structured output for critique via Pydantic model.
|
|
- Demo function sort_numbers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TypedDict, Dict
|
|
|
|
from langgraph.graph import StateGraph, END
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage
|
|
from pydantic import BaseModel, Field
|
|
|
|
# ---------- State ----------
|
|
class CodeReviewState(TypedDict):
|
|
code: str
|
|
draft_review: str
|
|
criteria_scores: Dict[str, int]
|
|
weakest_criterion: str
|
|
verdict: str
|
|
round: int
|
|
max_rounds: int
|
|
|
|
# ---------- LLM ----------
|
|
llm = ChatOpenAI(temperature=0)
|
|
|
|
# ---------- Nodes ----------
|
|
class DraftReviewOutput(BaseModel):
|
|
review: str = Field(..., description="Draft review text")
|
|
|
|
class ReflectOutput(BaseModel):
|
|
scores: Dict[str, int] = Field(..., description="Scores 0-10 for each criterion")
|
|
weakest_criterion: str = Field(..., description="Criterion with lowest score")
|
|
verdict: str = Field(..., description="'ok' or 'needs_revision'")
|
|
|
|
|
|
def draft_review(state: CodeReviewState) -> CodeReviewState:
|
|
code = state["code"]
|
|
prompt = f"""
|
|
Write a concise code review (3-6 bullet points) for the following Python function. Focus on style, correctness, and potential improvements.
|
|
|
|
```python
|
|
{code}
|
|
```
|
|
"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
data = DraftReviewOutput.model_validate_json(response.content)
|
|
state["draft_review"] = data.review
|
|
return state
|
|
|
|
|
|
def reflect(state: CodeReviewState) -> CodeReviewState:
|
|
review = state["draft_review"]
|
|
prompt = f"""
|
|
Evaluate the following review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10. Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
|
|
|
|
Review:
|
|
{review}
|
|
"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
data = ReflectOutput.model_validate_json(response.content)
|
|
state["criteria_scores"] = data.scores
|
|
state["weakest_criterion"] = data.weakest_criterion
|
|
state["verdict"] = data.verdict
|
|
return state
|
|
|
|
|
|
def rewrite(state: CodeReviewState) -> CodeReviewState:
|
|
crit = state["weakest_criterion"]
|
|
review = state["draft_review"]
|
|
prompt = f"""
|
|
The review below is weak in the {crit} criterion. Rewrite the review to improve that aspect.
|
|
Original review:
|
|
{review}
|
|
"""
|
|
new_review = llm.invoke([HumanMessage(content=prompt)])
|
|
state["draft_review"] = new_review.content
|
|
state["round"] += 1
|
|
return state
|
|
|
|
# ---------- Graph ----------
|
|
builder = StateGraph(CodeReviewState)
|
|
builder.add_node("draft_review", draft_review)
|
|
builder.add_node("reflect", reflect)
|
|
builder.add_node("rewrite", rewrite)
|
|
|
|
builder.set_entry_point("draft_review")
|
|
# After draft_review, go to reflect
|
|
builder.add_edge("draft_review", "reflect")
|
|
# After reflect, decide
|
|
builder.add_conditional_edges(
|
|
"reflect",
|
|
lambda state: state["verdict"] == "ok",
|
|
{"ok": END, "needs_revision": "rewrite"},
|
|
)
|
|
# After rewrite, go back to reflect if rounds remain
|
|
builder.add_conditional_edges(
|
|
"rewrite",
|
|
lambda state: state["round"] < state["max_rounds"],
|
|
{"continue": "reflect", "end": END},
|
|
)
|
|
# Compile graph
|
|
|
|
graph = builder.compile()
|
|
|
|
# ---------- Demo ----------
|
|
|
|
def sort_numbers(arr):
|
|
return sorted(arr)
|
|
|
|
if __name__ == "__main__":
|
|
code_str = "def sort_numbers(arr):\n return sorted(arr)\n"
|
|
initial_state: CodeReviewState = {
|
|
"code": code_str,
|
|
"draft_review": "",
|
|
"criteria_scores": {},
|
|
"weakest_criterion": "",
|
|
"verdict": "",
|
|
"round": 0,
|
|
"max_rounds": 2,
|
|
}
|
|
result = graph.invoke(initial_state)
|
|
print("Final state:")
|
|
print(result)
|