diff --git a/src/main.py b/src/main.py deleted file mode 100644 index 6b777ec..0000000 --- a/src/main.py +++ /dev/null @@ -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"])