Implemented LangGraph code review agent with LLM-based reflect and rewrite logic. Removed dummy src/main.py.: remove src/main.py
This commit is contained in:
-96
@@ -1,96 +0,0 @@
|
|||||||
"""
|
|
||||||
LangGraph code review agent with reflection on four criteria.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from typing import TypedDict, Dict
|
|
||||||
import os
|
|
||||||
|
|
||||||
# State definition
|
|
||||||
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
|
|
||||||
|
|
||||||
# Dummy LLM functions (replace with real LLM calls)
|
|
||||||
|
|
||||||
def draft_review(code: str) -> str:
|
|
||||||
return f"Review of code:\n- Function looks fine.\n- Consider adding type hints."
|
|
||||||
|
|
||||||
def reflect(review: str, code: str) -> Dict[str, int]:
|
|
||||||
# Simple heuristic scores
|
|
||||||
scores = {
|
|
||||||
"pep8": 7,
|
|
||||||
"type_hints": 5,
|
|
||||||
"edge_cases": 6,
|
|
||||||
"naming": 8,
|
|
||||||
}
|
|
||||||
weakest = min(scores, key=scores.get)
|
|
||||||
verdict = "needs_revision" if scores[weakest] < 7 else "ok"
|
|
||||||
return {**scores, "weakest_criterion": weakest, "verdict": verdict}
|
|
||||||
|
|
||||||
def rewrite(review: str, criterion: str) -> str:
|
|
||||||
return review + f"\n- Improve {criterion} section."
|
|
||||||
|
|
||||||
# Graph logic
|
|
||||||
from langgraph.graph import StateGraph
|
|
||||||
|
|
||||||
def build_graph() -> StateGraph[CodeReviewState]:
|
|
||||||
graph = StateGraph(CodeReviewState)
|
|
||||||
|
|
||||||
def start_node(state: CodeReviewState):
|
|
||||||
state["draft_review"] = draft_review(state["code"])
|
|
||||||
return state
|
|
||||||
|
|
||||||
def reflect_node(state: CodeReviewState):
|
|
||||||
scores = reflect(state["draft_review"], state["code"])
|
|
||||||
for k, v in scores.items():
|
|
||||||
if k in ["pep8", "type_hints", "edge_cases", "naming"]:
|
|
||||||
state["criteria_scores"][k] = v
|
|
||||||
state["weakest_criterion"] = scores["weakest_criterion"]
|
|
||||||
state["verdict"] = scores["verdict"]
|
|
||||||
return state
|
|
||||||
|
|
||||||
def rewrite_node(state: CodeReviewState):
|
|
||||||
state["draft_review"] = rewrite(state["draft_review"], state["weakest_criterion"])
|
|
||||||
state["round"] += 1
|
|
||||||
return state
|
|
||||||
|
|
||||||
graph.add_node("start", start_node)
|
|
||||||
graph.add_node("reflect", reflect_node)
|
|
||||||
graph.add_node("rewrite", rewrite_node)
|
|
||||||
|
|
||||||
graph.set_entry_point("start")
|
|
||||||
graph.add_edge("start", "reflect")
|
|
||||||
graph.add_conditional_edges(
|
|
||||||
"reflect",
|
|
||||||
lambda state: "end" if state["verdict"] == "ok" else "rewrite",
|
|
||||||
)
|
|
||||||
graph.add_edge("rewrite", "reflect")
|
|
||||||
|
|
||||||
# End condition based on max_rounds
|
|
||||||
def end_condition(state):
|
|
||||||
return state["round"] >= state["max_rounds"] or state["verdict"] == "ok"
|
|
||||||
|
|
||||||
graph.set_finish_condition(end_condition)
|
|
||||||
|
|
||||||
return graph
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
code_example = "def sort_numbers(arr):\n return sorted(arr)"
|
|
||||||
initial_state: CodeReviewState = {
|
|
||||||
"code": code_example,
|
|
||||||
"draft_review": "",
|
|
||||||
"criteria_scores": {},
|
|
||||||
"weakest_criterion": "",
|
|
||||||
"verdict": "",
|
|
||||||
"round": 0,
|
|
||||||
"max_rounds": 2,
|
|
||||||
}
|
|
||||||
graph = build_graph()
|
|
||||||
result = graph.invoke(initial_state)
|
|
||||||
print("Final draft review:\n", result["draft_review"])
|
|
||||||
print("Scores:", result["criteria_scores"])
|
|
||||||
Reference in New Issue
Block a user