add main.py

main.py

@@ -1,120 +1,203 @@
-# main.py
+"""LangGraph Code Review Agent with Reflection.
+
+This repository implements a LangGraph agent that takes a Python function as input and
+produces a code review. A critic node evaluates the review on four criteria:
+
+* PEP8 compliance
+* Type hints
+* Edge cases handling
+* Naming conventions
+
+If the critic returns `needs_revision`, the `rewrite` node rewrites the weakest
+criterion section of the review. The process repeats until the critic is satisfied
+or the maximum number of rounds is reached.
+
+The demo in ``__main__`` shows how to run the agent on a simple function.
+"""
+
+from __future__ import annotations
+
 import os
 from typing import TypedDict, Dict
-from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import create_chat_agent
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage
 
-# Define state
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_react_agent
+from pydantic import BaseModel, Field
+
+# ---------------------------------------------------------------------------
+# State definition
+# ---------------------------------------------------------------------------
 class CodeReviewState(TypedDict):
     code: str
-    draft_review: str
+    draft_review: str | None
-    criteria_scores: Dict[str, int]
+    criteria_scores: Dict[str, int] | None
-    weakest_criterion: str
+    weakest_criterion: str | None
-    verdict: str
+    verdict: str | None  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
-# LLM
+# ---------------------------------------------------------------------------
-llm = ChatOpenAI(temperature=0)
+# LLM configuration
+# ---------------------------------------------------------------------------
+# The OpenAI API key must be set in the environment variable OPENAI_API_KEY.
+# For local Ollama usage, replace ChatOpenAI with ChatOllama.
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
 
-# Draft review node
+# ---------------------------------------------------------------------------
+# Node: draft_review
+# ---------------------------------------------------------------------------
 async def draft_review(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""
+    """Generate an initial code review.
-You are a senior Python developer. Review the following code and provide a concise code review (3-6 bullet points) highlighting what is good and what can be improved.
 
-Code:
+    The review contains 3–6 bullet points describing what is good and what can be
-{state['code']}
+    improved.  The output is plain text.
-
+    """
-Review:
+    prompt = (
-"""
+        "You are a senior Python developer.\n"
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+        "Given the following function, write a concise code review (3–6 points).\n"
-    state['draft_review'] = response.content
+        "Focus on style, correctness, and potential improvements.\n"
+        "Return only the review text.\n\n"
+        f"Function:\n{state["code"]}\n"
+    )
+    review = await llm.ainvoke(prompt)
+    state["draft_review"] = review.content.strip()
     return state
 
-# Reflect node
+# ---------------------------------------------------------------------------
+# Node: reflect
+# ---------------------------------------------------------------------------
+class ReflectOutput(BaseModel):
+    pep8: int = Field(..., ge=0, le=10)
+    type_hints: int = Field(..., ge=0, le=10)
+    edge_cases: int = Field(..., ge=0, le=10)
+    naming: int = Field(..., ge=0, le=10)
+    verdict: str = Field(..., regex="^(ok|needs_revision)$")
+    weakest_criterion: str = Field(..., regex="^(pep8|type_hints|edge_cases|naming)$")
+
 async def reflect(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""
+    """Critic node that scores the draft review on four criteria.
-You are an AI critic evaluating a code review. Assign a score 0-10 for each of the following criteria based on the draft review:
-- pep8
-- type_hints
-- edge_cases
-- naming
 
-Provide a JSON object with keys "pep8", "type_hints", "edge_cases", "naming" and integer values.
+    The LLM returns a JSON object that matches ``ReflectOutput``.  The function
-Also determine the weakest criterion (the one with lowest score) and a verdict: "ok" if all scores >=7, otherwise "needs_revision".
+    parses the JSON and updates the state.
-
+    """
-Draft review:
+    prompt = (
-{state['draft_review']}
+        "You are a code quality critic.\n"
-
+        "Given the following code review, score it on the following criteria (0–10):\n"
-Output JSON:
+        "- PEP8 compliance\n"
-"""
+        "- Type hints usage\n"
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+        "- Edge cases handling\n"
-    import json
+        "- Naming conventions\n"
-    scores = json.loads(response.content)
+        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, verdict, weakest_criterion.\n"
-    state['criteria_scores'] = scores
+        "Verdict should be "ok" if all scores are >=7, otherwise "needs_revision".\n"
-    weakest = min(scores, key=scores.get)
+        "Weakest criterion is the one with the lowest score.\n\n"
-    state['weakest_criterion'] = weakest
+        f"Review:\n{state["draft_review"]}\n"
-    state['verdict'] = "ok" if all(v >= 7 for v in scores.values()) else "needs_revision"
+    )
+    result = await llm.ainvoke(prompt)
+    try:
+        data = ReflectOutput.model_validate_json(result.content)
+    except Exception as e:
+        # Fallback: if parsing fails, treat as needs_revision
+        data = ReflectOutput(
+            pep8=0,
+            type_hints=0,
+            edge_cases=0,
+            naming=0,
+            verdict="needs_revision",
+            weakest_criterion="pep8",
+        )
+    state["criteria_scores"] = data.model_dump(exclude="verdict,weakest_criterion")
+    state["weakest_criterion"] = data.weakest_criterion
+    state["verdict"] = data.verdict
     return state
 
-# Rewrite node
+# ---------------------------------------------------------------------------
+# Node: rewrite
+# ---------------------------------------------------------------------------
 async def rewrite(state: CodeReviewState) -> CodeReviewState:
-    crit = state['weakest_criterion']
+    """Rewrite the weakest part of the review.
-    prompt = f"""
-You are a senior Python developer. Rewrite the section of the code review that addresses the {crit} criterion, improving it. Keep the rest of the review unchanged.
 
-Original review:
+    The LLM is instructed to rewrite only the section that addresses the
-{state['draft_review']}
+    weakest criterion.  The new review replaces the old one.
-
+    """
-Rewrite only the part related to {crit}:
+    prompt = (
-"""
+        "You are a senior Python developer.\n"
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+        "Rewrite the part of the following code review that addresses the weakest criterion.\n"
-    # Replace the part in draft_review that mentions crit
+        "Keep the rest of the review unchanged.\n"
-    # For simplicity, just append the new part
+        "Return only the updated review text.\n\n"
-    state['draft_review'] = state['draft_review'] + "\n" + response.content
+        f"Weakest criterion: {state["weakest_criterion"]}\n"
-    state['round'] += 1
+        f"Current review:\n{state["draft_review"]}\n"
+    )
+    new_review = await llm.ainvoke(prompt)
+    state["draft_review"] = new_review.content.strip()
+    state["round"] += 1
     return state
 
-# Build graph
+# ---------------------------------------------------------------------------
-builder = StateGraph(CodeReviewState)
+# Graph construction
-builder.add_node("draft_review", draft_review)
+# ---------------------------------------------------------------------------
-builder.add_node("reflect", reflect)
+def build_graph() -> StateGraph[CodeReviewState]:
-builder.add_node("rewrite", rewrite)
+    graph = StateGraph(CodeReviewState)
+    graph.add_node("draft_review", draft_review)
+    graph.add_node("reflect", reflect)
+    graph.add_node("rewrite", rewrite)
 
-builder.set_entry_point("draft_review")
+    # Entry point
-builder.add_edge("draft_review", "reflect")
+    graph.set_entry_point("draft_review")
-builder.add_conditional_edges(
-    "reflect",
-    lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END",
-)
-builder.add_edge("rewrite", "reflect")
 
-graph = builder.compile()
+    # Transitions
+    graph.add_conditional_edges(
+        "draft_review",
+        lambda state: "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")
 
-# Demo function
+    return graph
-async def run_demo():
-    code = """
-# Example function to sort numbers
 
-def sort_numbers(arr):
+# ---------------------------------------------------------------------------
-    return sorted(arr)
+# Demo helper
-"""
+# ---------------------------------------------------------------------------
-    init_state: CodeReviewState = {
+async def run_demo(func) -> None:
+    """Run the graph on a single function and print the results."""
+    import inspect
+    import asyncio
+
+    code = inspect.getsource(func)
+    state: CodeReviewState = {
         "code": code,
-        "draft_review": "",
+        "draft_review": None,
-        "criteria_scores": {},
+        "criteria_scores": None,
-        "weakest_criterion": "",
+        "weakest_criterion": None,
-        "verdict": "",
+        "verdict": None,
         "round": 0,
         "max_rounds": 2,
     }
-    result = await graph.ainvoke(init_state)
+    graph = build_graph()
-    print("Final Review:\n", result["draft_review"])
+    final_state = await graph.astate(state)
-    print("Scores:\n", result["criteria_scores"])
+    print("\n=== Initial Review ===")
-    print("Verdict:\n", result["verdict"])
+    print(state["draft_review"])
+    print("\n=== Scores ===")
+    print(final_state["criteria_scores"])
+    print("Verdict:", final_state["verdict"])
+    if final_state["verdict"] == "needs_revision":
+        print("\n=== Revised Review ===")
+        print(final_state["draft_review"])
 
+# ---------------------------------------------------------------------------
+# Example function for demo
+# ---------------------------------------------------------------------------
+
+def sort_numbers(arr):
+    """Return a sorted copy of the input list."""
+    return sorted(arr)
+
+# ---------------------------------------------------------------------------
+# Main entry point
+# ---------------------------------------------------------------------------
 if __name__ == "__main__":
     import asyncio
-    asyncio.run(run_demo())
+    asyncio.run(run_demo(sort_numbers))