118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
"""LangGraph code review agent.
|
|
|
|
Implementation follows assignment:
|
|
- State: CodeReviewState with 4 criteria.
|
|
- Nodes: review_and_critique, rewrite.
|
|
- Graph: START -> review_and_critique -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> review_and_critique).
|
|
- Uses LangGraph and LangChain OpenAI for LLM calls.
|
|
- Structured output for critique via Pydantic model.
|
|
- Demo function sort_numbers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
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 ----------
|
|
# Use OpenAI only
|
|
llm = ChatOpenAI(temperature=0)
|
|
|
|
# ---------- Nodes ----------
|
|
class CritiqueOutput(BaseModel):
|
|
review: str = Field(..., description="Draft review text")
|
|
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 review_and_critique(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}
|
|
```
|
|
|
|
Then evaluate the review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10.
|
|
Return a JSON object with keys: review (string), scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
|
|
"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
data = CritiqueOutput.model_validate_json(response.content)
|
|
state["draft_review"] = data.review
|
|
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("review_and_critique", review_and_critique)
|
|
builder.add_node("rewrite", rewrite)
|
|
|
|
builder.set_entry_point("review_and_critique")
|
|
# After initial review_and_critique, decide to end if verdict ok
|
|
builder.add_conditional_edges(
|
|
"review_and_critique",
|
|
lambda state: state["verdict"] == "ok",
|
|
{"ok": END, "needs_revision": "rewrite"},
|
|
)
|
|
# After rewrite, go back to review_and_critique if rounds remain
|
|
builder.add_conditional_edges(
|
|
"rewrite",
|
|
lambda state: "review_and_critique" if state["round"] < state["max_rounds"] else END,
|
|
{"review_and_critique": "review_and_critique", END: END}
|
|
)
|
|
|
|
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)
|