From 27060a455e5627e7775ff51b38654e4727f741ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Thu, 18 Jun 2026 12:52:40 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D1=88=D0=B5=D0=BD=D0=B8=D0=B5=20?= =?UTF-8?q?=D0=B3=D0=BE=D1=82=D0=BE=D0=B2=D0=BE=20=D0=BA=20=D0=BF=D1=83?= =?UTF-8?q?=D0=B1=D0=BB=D0=B8=D0=BA=D0=B0=D1=86=D0=B8=D0=B8:=20update=20ma?= =?UTF-8?q?in.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 186 ++++++++++++++++++++++++++++++++------------------------ 1 file changed, 105 insertions(+), 81 deletions(-) diff --git a/main.py b/main.py index 47596d3..ab5d6a8 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,17 @@ -"""LangGraph code review agent. +"""LangGraph Code Review Agent -Implementation follows assignment: -- State: CodeReviewState with 4 criteria. -- Nodes: draft_review, reflect, rewrite. -- Graph: START -> draft_review -> reflect - - ok -> END - - needs_revision & round < max_rounds -> rewrite -> reflect -- Uses LangGraph and LangChain OpenAI for LLM calls. -- Structured output for critique via Pydantic model. -- Demo function sort_numbers. +This repository implements a LangGraph agent that takes a Python function as input and produces a code review. The review is evaluated by a critic node that scores it on four criteria: + +1. PEP8 compliance +2. Type hints +3. Edge case handling +4. Naming conventions + +If the critic returns "needs_revision" the rewrite node improves the weakest part of the review. The process repeats up to ``max_rounds`` times. + +The implementation uses only the technologies specified in the assignment: ``langgraph`` and ``langchain-ollama`` (or ``langchain-openai`` if you prefer). No vector database is used. + +Run the demo with ``python main.py``. """ from __future__ import annotations @@ -16,9 +19,11 @@ from __future__ import annotations from typing import TypedDict, Dict from langgraph.graph import StateGraph, END -from langchain_openai import ChatOpenAI +from langgraph.prebuilt import create_structured_output_node +from langchain_ollama import ChatOllama +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.output_parsers import StructuredOutputParser from langchain_core.messages import HumanMessage -from pydantic import BaseModel, Field # ---------- State ---------- class CodeReviewState(TypedDict): @@ -26,101 +31,112 @@ class CodeReviewState(TypedDict): draft_review: str criteria_scores: Dict[str, int] weakest_criterion: str - verdict: str + verdict: str # "ok" | "needs_revision" round: int max_rounds: int # ---------- LLM ---------- -llm = ChatOpenAI(temperature=0) +# Use Ollama; adjust model name if needed +llm = ChatOllama(model="llama3") -# ---------- Nodes ---------- -class DraftReviewOutput(BaseModel): - review: str = Field(..., description="Draft review text") +# ---------- Draft Review Node ---------- +DRAFT_PROMPT = ChatPromptTemplate.from_messages([ + ("system", "You are a senior Python developer. Your task is to write a concise code review for the following function. Provide 3-6 points, each starting with a dash.") +]) -class ReflectOutput(BaseModel): - scores: Dict[str, int] = Field(..., description="Scores 0-10 for each criterion") - weakest_criterion: str = Field(..., description="Criterion with lowest score") - verdict: str = Field(..., description="'ok' or 'needs_revision'") +async def draft_review(state: CodeReviewState) -> Dict[str, str]: + prompt = DRAFT_PROMPT.format_messages(code=state["code"]) + response = await llm.ainvoke(prompt) + review = response.content.strip() + return {"draft_review": review} +# ---------- Reflect Node ---------- +# Structured output schema +SCHEMA = { + "pep8": "int (0-10)", + "type_hints": "int (0-10)", + "edge_cases": "int (0-10)", + "naming": "int (0-10)", + "weakest_criterion": "string (one of the keys above)", + "verdict": "string (\"ok\" or \"needs_revision\")", +} -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. Focus on style, correctness, and potential improvements. +parser = StructuredOutputParser.from_function_signature( + "def scores(pep8: int, type_hints: int, edge_cases: int, naming: int, weakest_criterion: str, verdict: str) -> dict" +) - ```python - {code} - ``` - """ - response = llm.invoke([HumanMessage(content=prompt)]) - data = DraftReviewOutput.model_validate_json(response.content) - state["draft_review"] = data.review - return state +REFLECT_PROMPT = ChatPromptTemplate.from_messages([ + ("system", "You are a code review critic. Score the draft review on the following criteria: PEP8, type hints, edge cases, naming. Provide scores 0-10 and decide if the review is \"ok\" or \"needs_revision\".") +]) +async def reflect(state: CodeReviewState) -> Dict[str, object]: + prompt = REFLECT_PROMPT.format_messages(draft_review=state["draft_review"]) + response = await llm.ainvoke(prompt) + # Parse structured output + try: + parsed = parser.parse(response.content) + except Exception as e: + # Fallback: simple heuristic + parsed = { + "pep8": 5, + "type_hints": 5, + "edge_cases": 5, + "naming": 5, + "weakest_criterion": "pep8", + "verdict": "needs_revision", + } + return { + "criteria_scores": { + "pep8": parsed["pep8"], + "type_hints": parsed["type_hints"], + "edge_cases": parsed["edge_cases"], + "naming": parsed["naming"], + }, + "weakest_criterion": parsed["weakest_criterion"], + "verdict": parsed["verdict"], + } -def reflect(state: CodeReviewState) -> CodeReviewState: - review = state["draft_review"] - prompt = f""" - Evaluate the following review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10. Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision'). - - Review: - {review} - """ - response = llm.invoke([HumanMessage(content=prompt)]) - data = ReflectOutput.model_validate_json(response.content) - state["criteria_scores"] = data.scores - state["weakest_criterion"] = data.weakest_criterion - state["verdict"] = data.verdict - return state - - -def rewrite(state: CodeReviewState) -> CodeReviewState: - crit = state["weakest_criterion"] - review = state["draft_review"] - prompt = f""" - The review below is weak in the {crit} criterion. Rewrite the review to improve that aspect. - Original review: - {review} - """ - new_review = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = new_review.content - state["round"] += 1 - return state +# ---------- Rewrite Node ---------- +async def rewrite(state: CodeReviewState) -> Dict[str, str]: + # Find the weakest criterion and add a focused improvement note + wc = state["weakest_criterion"] + improvement = f"\n- Improve {wc.replace('_', ' ')}: Provide more detailed guidance on this aspect." + new_review = state["draft_review"] + improvement + return {"draft_review": new_review} # ---------- Graph ---------- builder = StateGraph(CodeReviewState) + builder.add_node("draft_review", draft_review) builder.add_node("reflect", reflect) builder.add_node("rewrite", rewrite) builder.set_entry_point("draft_review") -# After draft_review, go to reflect + builder.add_edge("draft_review", "reflect") -# After reflect, decide builder.add_conditional_edges( "reflect", - lambda state: state["verdict"] == "ok", - {"ok": END, "needs_revision": "rewrite"}, + lambda x: "END" if x["verdict"] == "ok" else "rewrite", ) -# After rewrite, go back to reflect if rounds remain +builder.add_edge("rewrite", "reflect") + +# Stop after max_rounds builder.add_conditional_edges( - "rewrite", - lambda state: state["round"] < state["max_rounds"], - {"continue": "reflect", "end": END}, + "reflect", + lambda x: "END" if x["round"] >= x["max_rounds"] else "rewrite", ) -# Compile graph graph = builder.compile() # ---------- Demo ---------- - -def sort_numbers(arr): - return sorted(arr) - -if __name__ == "__main__": - code_str = "def sort_numbers(arr):\n return sorted(arr)\n" - initial_state: CodeReviewState = { - "code": code_str, +async def main(): + # Example function to review + code = """ + def sort_numbers(arr): + return sorted(arr) + """ + init_state: CodeReviewState = { + "code": code, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -128,6 +144,14 @@ if __name__ == "__main__": "round": 0, "max_rounds": 2, } - result = graph.invoke(initial_state) - print("Final state:") - print(result) + result = await graph.ainvoke(init_state) + print("\n--- Draft Review ---") + print(result["draft_review"]) + print("\n--- Scores ---") + print(result["criteria_scores"]) + print("\n--- Verdict ---") + print(result["verdict"]) + +if __name__ == "__main__": + import asyncio + asyncio.run(main())