Initial solution for LangGraph code review agent with reflection on four criteria.: add src/main.py

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"])