From 4644f31171c3a9bb22be1e8b05d14bc5ec61cbd0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Fri, 3 Jul 2026 11:46:56 +0000 Subject: [PATCH] Initial commit of LangGraph code review agent --- main.py | 272 ++++++++++++++++++++++++++++++++++---------------------- 1 file changed, 167 insertions(+), 105 deletions(-) diff --git a/main.py b/main.py index a7e185f..71f0eaf 100644 --- a/main.py +++ b/main.py @@ -1,132 +1,173 @@ +""" +# main.py +# LangGraph code review agent with deepagents integration +# Author: OpenAI ChatGPT +# Requirements: deepagents, langchain>=1.2.10, langchain-openai>=0.3.0, langgraph>=0.2.0 +# Run: python main.py +""" + import os +import json import asyncio -from typing import TypedDict, Dict, Annotated +from typing import TypedDict, Annotated, Dict + from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.messages import HumanMessage from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langgraph.graph import StateGraph, START, END, add_conditional_edges +from langgraph.graph import StateGraph, START, END +from langgraph.graph.message import add_messages from pydantic import BaseModel, Field -from langchain_core.output_parsers import PydanticOutputParser -# DESIGN DECISION: Add deepagents to requirements.txt -# NECESSITY: deepagents is required for create_deep_agent usage -# OPTIMALITY: ensures reproducible installation -# ALTERNATIVES CONSIDERED: manual installation or alternative agent framework, but violates course requirement +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- +# Load OpenRouter API key from environment +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY environment variable not set") -# Load environment variables -from dotenv import load_dotenv -load_dotenv() - -# LLM configuration using OpenRouter +# LLM instance (OpenRouter) llm = ChatOpenAI( - model="openai/gpt-oss-20b", + model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), + api_key=OPENAI_API_KEY, temperature=0.0, ) -# Backend for deepagents (not heavily used but required by create_deep_agent) -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# Deepagents agent used for drafting and rewriting reviews -agent = create_deep_agent( - model=llm, - tools=[], - backend=backend, - system_prompt="You are a helpful agent.", -) - +# --------------------------------------------------------------------------- # State definition +# --------------------------------------------------------------------------- class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: Dict[str, int] + criteria_scores: Dict[str, int] # {"pep8": int, "type_hints": int, "edge_cases": int, "naming": int} weakest_criterion: str verdict: str # "ok" | "needs_revision" round: int max_rounds: int -# Structured output for the reflect node -class CritiqueOutput(BaseModel): - scores: Dict[str, int] = Field(description="Scores 0-10 for each criterion") - verdict: str = Field(description="ok or needs_revision") - weakest_criterion: str = Field(description="Criterion with lowest score") +# --------------------------------------------------------------------------- +# Pydantic model for reflect output +# --------------------------------------------------------------------------- +class ReviewScores(BaseModel): + pep8: int = Field(..., description="Score for PEP8 compliance (0-10)") + type_hints: int = Field(..., description="Score for type hints usage (0-10)") + edge_cases: int = Field(..., description="Score for handling edge cases (0-10)") + naming: int = Field(..., description="Score for naming conventions (0-10)") + verdict: str = Field(..., description="'ok' if all scores >=7 else 'needs_revision'") -parser = PydanticOutputParser(pydantic_object=CritiqueOutput) +# --------------------------------------------------------------------------- +# LangGraph nodes +# --------------------------------------------------------------------------- +async def draft_review(state: CodeReviewState) -> CodeReviewState: + code = state["code"] + prompt = f""" +Write a concise code review (3-6 bullet points) for the following Python function: -# Node: draft_review -async def draft_review_node(state: CodeReviewState) -> CodeReviewState: - system_msg = SystemMessage(content="You are a code reviewer. Provide a concise review (3-6 points) of the following Python function.") - user_msg = HumanMessage(content=state["code"]) - result = await agent.ainvoke( - {"messages": [system_msg, user_msg]}, - {"configurable": {"thread_id": "draft-review"}}, - ) - review = result["messages"][-1].content +{code} + +Review: +""" + messages = [HumanMessage(content=prompt)] + response = await llm.ainvoke(messages) + review = response.content.strip() state["draft_review"] = review return state -# Node: reflect -async def reflect_node(state: CodeReviewState) -> CodeReviewState: - system_msg = SystemMessage(content="You are a code critic. Score the following review on four criteria: pep8, type_hints, edge_cases, naming. Return a JSON with scores, verdict, and weakest_criterion.") - user_msg = HumanMessage(content=state["draft_review"]) - raw_output = await llm.invoke([system_msg, user_msg]) - critique = parser.parse(raw_output.content) - state["criteria_scores"] = critique.scores - state["weakest_criterion"] = critique.weakest_criterion - state["verdict"] = critique.verdict +async def reflect(state: CodeReviewState) -> CodeReviewState: + review = state["draft_review"] + code = state["code"] + prompt = f""" +You are a code review critic. Evaluate the following review against the code. Assign scores 0-10 for each criterion: PEP8, type hints, edge cases, naming. Also provide verdict: "ok" if all scores >=7, else "needs_revision". Return JSON with keys: pep8, type_hints, edge_cases, naming, verdict. + +Example: +{{"pep8":8,"type_hints":9,"edge_cases":6,"naming":7,"verdict":"needs_revision"}} + +Code: +{code} + +Review: +{review} + +JSON: +""" + messages = [HumanMessage(content=prompt)] + response = await llm.ainvoke(messages) + try: + data = json.loads(response.content) + except Exception as e: + data = {} + scores = { + "pep8": int(data.get("pep8", 0)), + "type_hints": int(data.get("type_hints", 0)), + "edge_cases": int(data.get("edge_cases", 0)), + "naming": int(data.get("naming", 0)), + } + weakest = min(scores, key=scores.get) + state["criteria_scores"] = scores + state["weakest_criterion"] = weakest + state["verdict"] = data.get("verdict", "needs_revision") return state -# Node: rewrite -async def rewrite_node(state: CodeReviewState) -> CodeReviewState: - criterion = state["weakest_criterion"] - system_msg = SystemMessage(content=f"You are a code reviewer. Rewrite the review to improve the section about {criterion}. Keep the overall structure.") - user_msg = HumanMessage(content=state["draft_review"]) - result = await agent.ainvoke( - {"messages": [system_msg, user_msg]}, - {"configurable": {"thread_id": "rewrite"}}, - ) - new_review = result["messages"][-1].content - state["draft_review"] = new_review +async def rewrite(state: CodeReviewState) -> CodeReviewState: + review = state["draft_review"] + weakest = state["weakest_criterion"] + code = state["code"] + prompt = f""" +Rewrite the part of the review that addresses the {weakest} criterion to improve it. Keep other parts unchanged. Provide only the updated review. + +Original review: +{review} + +Updated review: +""" + messages = [HumanMessage(content=prompt)] + response = await llm.ainvoke(messages) + updated_review = response.content.strip() + state["draft_review"] = updated_review state["round"] += 1 return state -# Conditional edge function -def decide_next(state: CodeReviewState) -> str: - if state["verdict"] == "ok": - return "END" - if state["round"] < state["max_rounds"]: +# --------------------------------------------------------------------------- +# Build the graph +# --------------------------------------------------------------------------- +graph = StateGraph(CodeReviewState) +graph.add_node("draft_review", draft_review) +graph.add_node("reflect", reflect) +graph.add_node("rewrite", rewrite) + +graph.set_entry_point("draft_review") +# After draft_review always go to reflect +graph.add_edge("draft_review", "reflect") +# Conditional after reflect +def reflect_cond(state: CodeReviewState): + if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]: return "rewrite" return "END" -# Build the graph -graph = StateGraph(CodeReviewState) -graph.add_node("draft_review", draft_review_node) -graph.add_node("reflect", reflect_node) -graph.add_node("rewrite", rewrite_node) -graph.add_conditional_edges("reflect", decide_next, { - "rewrite": "rewrite", - "END": END, -}) -graph.add_edge(START, "draft_review") -graph.add_edge("draft_review", "reflect") +graph.add_conditional_edges("reflect", reflect_cond) +# After rewrite go back to reflect graph.add_edge("rewrite", "reflect") -app = graph.compile() -# Demo function -def demo_code() -> str: - return """def sort_numbers(arr): - return sorted(arr)""" +compiled_graph = graph.compile() -async def main(): - code_str = demo_code() - initial_state: CodeReviewState = { - "code": code_str, +# --------------------------------------------------------------------------- +# DeepAgents integration +# --------------------------------------------------------------------------- +# Backend for deepagents (in-memory shell, no real files used) +backend = CompositeBackend( + default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), + routes={}, +) + +@tool +def run_code_review(code: str) -> str: + """Run a multi‑round code review on the provided Python function.""" + # Initial state + state: CodeReviewState = { + "code": code, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -134,20 +175,41 @@ async def main(): "round": 0, "max_rounds": 2, } - final_state = await app.ainvoke(initial_state) - print("\n=== Initial Draft Review ===") - print(initial_state["draft_review"]) - print("\n=== Scores ===") - print(final_state["criteria_scores"]) - print("\n=== Weakest Criterion ===") - print(final_state["weakest_criterion"]) - if final_state["verdict"] == "needs_revision": - print("\n=== Revised Review ===") - print(final_state["draft_review"]) - print("\n=== Updated Scores ===") - print(final_state["criteria_scores"]) - else: - print("\nReview is satisfactory. No rewrite needed.") + async def _run(): + final_state = await compiled_graph.invoke(state) + # Logging + print("\n--- Draft Review ---") + print(final_state["draft_review"]) + print("\n--- Scores ---") + print(final_state["criteria_scores"]) + if final_state["verdict"] == "needs_revision": + print("\n--- Rewrite performed ---") + print("\n--- Final Review ---") + print(final_state["draft_review"]) + return final_state["draft_review"] + + return asyncio.run(_run()) + +agent = create_deep_agent( + model=llm, + tools=[run_code_review], + backend=backend, + system_prompt="You are a helpful code review assistant. Use the provided tool to review code.", +) + +# --------------------------------------------------------------------------- +# Demo CLI +# --------------------------------------------------------------------------- if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + sample_code = """ +# Example function to sort numbers + +def sort_numbers(arr): + return sorted(arr) +""" + # Directly invoke the tool for demonstration + print("Running demo on sample function...") + final_review = run_code_review(sample_code) + print("\n=== Final Review Output ===") + print(final_review)