From faaa1d85b6bf5d1d38af993f58a1cc080172cb8c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Thu, 11 Jun 2026 16:09:36 +0000 Subject: [PATCH] add main.py --- main.py | 240 +++++++++++++++++++------------------------------------- 1 file changed, 81 insertions(+), 159 deletions(-) diff --git a/main.py b/main.py index 12cba81..e2eb1a9 100644 --- a/main.py +++ b/main.py @@ -1,198 +1,120 @@ -""" -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 - +# main.py import os from typing import TypedDict, Dict - from langgraph.graph import StateGraph, END -from langgraph.prebuilt import create_react_agent +from langgraph.prebuilt import create_chat_agent from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage - -# --------------------------------------------------------------------------- -# State definition -# --------------------------------------------------------------------------- +from langchain_core.messages import HumanMessage, AIMessage +# Define state class CodeReviewState(TypedDict): code: str draft_review: str criteria_scores: Dict[str, int] weakest_criterion: str - verdict: str # "ok" | "needs_revision" + verdict: str 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 -# --------------------------------------------------------------------------- +# LLM +llm = ChatOpenAI(temperature=0) +# Draft review node async def draft_review(state: CodeReviewState) -> CodeReviewState: - """Generate an initial review of the provided code. + prompt = f""" +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. - 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() +Code: +{state['code']} + +Review: +""" + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = response.content return state +# Reflect node async def reflect(state: CodeReviewState) -> CodeReviewState: - """Critic that scores the draft review on four criteria. + prompt = f""" +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 - 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." - ) +Provide a JSON object with keys "pep8", "type_hints", "edge_cases", "naming" and integer values. +Also determine the weakest criterion (the one with lowest score) and a verdict: "ok" if all scores >=7, otherwise "needs_revision". + +Draft review: +{state['draft_review']} + +Output JSON: +""" 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") + scores = json.loads(response.content) + state['criteria_scores'] = scores + weakest = min(scores, key=scores.get) + state['weakest_criterion'] = weakest + state['verdict'] = "ok" if all(v >= 7 for v in scores.values()) else "needs_revision" return state +# Rewrite node async def rewrite(state: CodeReviewState) -> CodeReviewState: - """Rewrite the weakest part of the review. + crit = state['weakest_criterion'] + 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. - 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 +Original review: +{state['draft_review']} + +Rewrite only the part related to {crit}: +""" + response = await llm.ainvoke([HumanMessage(content=prompt)]) + # Replace the part in draft_review that mentions crit + # For simplicity, just append the new part + state['draft_review'] = state['draft_review'] + "\n" + response.content + state['round'] += 1 return state -# --------------------------------------------------------------------------- -# Graph construction -# --------------------------------------------------------------------------- +# Build graph +builder = StateGraph(CodeReviewState) +builder.add_node("draft_review", draft_review) +builder.add_node("reflect", reflect) +builder.add_node("rewrite", rewrite) -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) +builder.set_entry_point("draft_review") +builder.add_edge("draft_review", "reflect") +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") - # Entry point - graph.set_entry_point("draft_review") +graph = builder.compile() - # 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") +# Demo function +async def run_demo(): + code = """ +# Example function to sort numbers - 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, +def sort_numbers(arr): + return sorted(arr) +""" + init_state: CodeReviewState = { + "code": code, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", "verdict": "", "round": 0, - "max_rounds": args.max_rounds, + "max_rounds": 2, } + result = await graph.ainvoke(init_state) + print("Final Review:\n", result["draft_review"]) + print("Scores:\n", result["criteria_scores"]) + print("Verdict:\n", result["verdict"]) - 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']}") +if __name__ == "__main__": + import asyncio + asyncio.run(run_demo())