add main.py

main.py

@@ -0,0 +1,198 @@
+"""
+LangGraph Code Review Agent
+===========================
+
+This repository contains a minimal LangGraph implementation that
+performs a code review on a Python function.  The graph consists of
+three nodes:
+
+* ``draft_review`` – generates an initial review.
+* ``reflect`` – a critic that scores the review on four criteria
+  (PEP8, type hints, edge cases, naming) and decides whether a
+  rewrite is required.
+* ``rewrite`` – rewrites the weakest part of the review.
+
+The graph runs for a maximum of ``max_rounds`` (default 2).  The
+demo can be executed with ``python -m main``.
+"""
+
+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
+
+# ---------------------------------------------------------------------------
+# 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
+
+# ---------------------------------------------------------------------------
+# LLM configuration
+# ---------------------------------------------------------------------------
+
+# The user must set the OPENAI_API_KEY environment variable.
+llm = ChatOpenAI(temperature=0.0, model="gpt-4o-mini")
+
+# ---------------------------------------------------------------------------
+# Node implementations
+# ---------------------------------------------------------------------------
+
+async def draft_review(state: CodeReviewState) -> CodeReviewState:
+    """Generate an initial review of the provided code.
+
+    The review is a short list of 3–6 bullet points.
+    """
+    code = state["code"]
+    prompt = (
+        "You are a senior Python developer.\n"
+        "Review the following function and provide a concise list of 3–6 points\n"
+        "highlighting what is good and what could be improved.\n"
+        "Do not mention the criteria – just give the review.\n"
+        f"Function:\n{code}\n"
+        "Review:"  # LLM will continue after this
+    )
+    review = await llm.ainvoke([HumanMessage(content=prompt)])
+    state["draft_review"] = review.content.strip()
+    return state
+
+async def reflect(state: CodeReviewState) -> CodeReviewState:
+    """Critic that scores the draft review on four criteria.
+
+    The output is a JSON object with keys:
+    * pep8, type_hints, edge_cases, naming – integers 0‑10
+    * weakest_criterion – one of the four keys
+    * verdict – "ok" or "needs_revision"
+    """
+    review = state["draft_review"]
+    code = state["code"]
+    prompt = (
+        "You are a code quality critic.\n"
+        "Given the following code and its draft review, score the review on\n"
+        "four criteria: PEP8, type hints, edge cases, naming.\n"
+        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming,\n"
+        "weakest_criterion, verdict.\n"
+        f"Code:\n{code}\n"
+        f"Draft review:\n{review}\n"
+        "Answer in JSON only."
+    )
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    import json
+    try:
+        scores = json.loads(response.content)
+    except Exception as e:
+        # Fallback: if parsing fails, treat as needs_revision
+        scores = {
+            "pep8": 0,
+            "type_hints": 0,
+            "edge_cases": 0,
+            "naming": 0,
+            "weakest_criterion": "pep8",
+            "verdict": "needs_revision",
+        }
+    state["criteria_scores"] = {
+        "pep8": int(scores.get("pep8", 0)),
+        "type_hints": int(scores.get("type_hints", 0)),
+        "edge_cases": int(scores.get("edge_cases", 0)),
+        "naming": int(scores.get("naming", 0)),
+    }
+    state["weakest_criterion"] = scores.get("weakest_criterion", "pep8")
+    state["verdict"] = scores.get("verdict", "needs_revision")
+    return state
+
+async def rewrite(state: CodeReviewState) -> CodeReviewState:
+    """Rewrite the weakest part of the review.
+
+    The node receives the current state and the weakest criterion.
+    It generates a new review that specifically addresses that criterion.
+    """
+    code = state["code"]
+    weakest = state["weakest_criterion"]
+    prompt = (
+        "You are a senior Python developer.\n"
+        "Rewrite the draft review to improve the part related to the following criterion: "
+        f"{weakest}.\n"
+        "Keep the rest of the review unchanged.\n"
+        "Output only the updated review.\n"
+        f"Current draft review:\n{state['draft_review']}\n"
+    )
+    new_review = await llm.ainvoke([HumanMessage(content=prompt)])
+    state["draft_review"] = new_review.content.strip()
+    state["round"] += 1
+    return state
+
+# ---------------------------------------------------------------------------
+# Graph construction
+# ---------------------------------------------------------------------------
+
+def build_graph() -> StateGraph[CodeReviewState]:
+    graph = StateGraph(CodeReviewState)
+    graph.add_node("draft_review", draft_review)
+    graph.add_node("reflect", reflect)
+    graph.add_node("rewrite", rewrite)
+
+    # Entry point
+    graph.set_entry_point("draft_review")
+
+    # Transitions
+    graph.add_edge("draft_review", "reflect")
+    graph.add_conditional_edges(
+        "reflect",
+        lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END",
+    )
+    graph.add_edge("rewrite", "reflect")
+
+    return graph
+
+# ---------------------------------------------------------------------------
+# Demo execution
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    import textwrap
+
+    parser = argparse.ArgumentParser(description="Run the code review graph on a demo function.")
+    parser.add_argument("--max-rounds", type=int, default=2, help="Maximum number of rewrite rounds")
+    args = parser.parse_args()
+
+    # Demo function – can be replaced by any user code
+    demo_code = textwrap.dedent(
+        """
+        def sort_numbers(arr):
+            return sorted(arr)
+        """
+    ).strip()
+
+    initial_state: CodeReviewState = {
+        "code": demo_code,
+        "draft_review": "",
+        "criteria_scores": {},
+        "weakest_criterion": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": args.max_rounds,
+    }
+
+    graph = build_graph()
+    final_state = graph.invoke(initial_state)
+
+    print("\n=== Final Review ===")
+    print(final_state["draft_review"])
+    print("\n=== Scores ===")
+    for k, v in final_state["criteria_scores"].items():
+        print(f"{k}: {v}")
+    print(f"Verdict: {final_state['verdict']}")
+    print(f"Rounds: {final_state['round']}")