142 lines
4.4 KiB
Python
142 lines
4.4 KiB
Python
"""LangGraph code review agent.
|
|
|
|
This implementation follows the assignment specification:
|
|
- State: CodeReviewState with 4 criteria.
|
|
- Nodes: draft_review, reflect, rewrite.
|
|
- Graph: START -> draft_review -> reflect -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> reflect).
|
|
- Uses LangGraph and LangChain OpenAI (or Ollama) for LLM calls.
|
|
- Structured output for reflect 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 langgraph.prebuilt import create_react_agent
|
|
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 if key present, else Ollama fallback
|
|
if os.getenv("OPENAI_API_KEY"):
|
|
llm = ChatOpenAI(temperature=0)
|
|
else:
|
|
llm = ChatOpenAI(model="ollama/llama3", temperature=0)
|
|
|
|
# ---------- Nodes ----------
|
|
|
|
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}
|
|
```
|
|
"""
|
|
review = llm.invoke([HumanMessage(content=prompt)])
|
|
state["draft_review"] = review.content
|
|
return state
|
|
|
|
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 reflect(state: CodeReviewState) -> CodeReviewState:
|
|
code = state["code"]
|
|
review = state["draft_review"]
|
|
prompt = f"""
|
|
You are a code review critic. Evaluate the following review of a Python function.
|
|
Function code:
|
|
```python
|
|
{code}
|
|
```
|
|
Review:
|
|
{review}
|
|
|
|
Score the 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').
|
|
"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
try:
|
|
data = ReflectOutput.model_validate_json(response.content)
|
|
except Exception:
|
|
# Fallback simple parsing
|
|
data = ReflectOutput.model_validate_json("{\"scores\":{\"pep8\":5,\"type_hints\":5,\"edge_cases\":5,\"naming\":5},\"weakest_criterion\":\"pep8\",\"verdict\":\"needs_revision\"}")
|
|
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")
|
|
builder.add_edge("draft_review", "reflect")
|
|
builder.add_conditional_edges(
|
|
"reflect",
|
|
lambda state: state["verdict"] == "ok",
|
|
{"ok": END, "needs_revision": "rewrite"},
|
|
)
|
|
# After rewrite, decide to reflect again or end if max rounds reached
|
|
builder.add_conditional_edges(
|
|
"rewrite",
|
|
lambda state: "reflect" if state["round"] < state["max_rounds"] else END,
|
|
{"reflect": "reflect", "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)
|