From 2a107613e295f9fe7d949ae3ef8d01fb2f7b3e6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Thu, 11 Jun 2026 15:36:49 +0000 Subject: [PATCH] Initial solution for LangGraph code review agent with reflection on four criteria.: add src/main.py --- src/main.py | 96 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 src/main.py diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..6b777ec --- /dev/null +++ b/src/main.py @@ -0,0 +1,96 @@ +""" +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"])