diff --git a/main.py b/main.py index ab5d6a8..b0720e4 100644 --- a/main.py +++ b/main.py @@ -1,31 +1,12 @@ -"""LangGraph Code Review Agent - -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 - from typing import TypedDict, Dict +import inspect from langgraph.graph import StateGraph, END -from langgraph.prebuilt import create_structured_output_node -from langchain_ollama import ChatOllama +from langchain_ollama import Ollama from langchain_core.prompts import ChatPromptTemplate -from langchain_core.output_parsers import StructuredOutputParser -from langchain_core.messages import HumanMessage +from langchain_core.output_parsers import JsonOutputParser -# ---------- State ---------- +# Define the state class CodeReviewState(TypedDict): code: str draft_review: str @@ -35,107 +16,81 @@ class CodeReviewState(TypedDict): round: int max_rounds: int -# ---------- LLM ---------- -# Use Ollama; adjust model name if needed -llm = ChatOllama(model="llama3") +# LLM instance (Ollama) +llm = Ollama(model="llama3.1") -# ---------- 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.") -]) +# Node: draft_review -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() +def draft_review(state: CodeReviewState) -> Dict[str, str]: + prompt = ChatPromptTemplate.from_messages([ + ("system", "You are a senior Python developer. Write a concise code review for the given function. Provide 3-6 actionable points."), + ("user", "Here is the function:\n{code}") + ]) + chain = prompt | llm + review = chain.invoke({"code": state["code"]}) 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\")", -} +# Node: reflect -parser = StructuredOutputParser.from_function_signature( - "def scores(pep8: int, type_hints: int, edge_cases: int, naming: int, weakest_criterion: str, verdict: str) -> dict" -) - -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", - } +def reflect(state: CodeReviewState) -> Dict[str, object]: + prompt = ChatPromptTemplate.from_messages([ + ("system", """You are a code quality critic. Score the following review on four criteria: PEP8, type hints, edge cases, naming. Return a JSON with integer scores 0-10, the weakest criterion, and verdict \"ok\" or \"needs_revision\".\n""") , + ("user", "Review:\n{draft_review}") + ]) + parser = JsonOutputParser() + chain = prompt | llm | parser + result = chain.invoke({"draft_review": state["draft_review"]}) + # result is a dict return { "criteria_scores": { - "pep8": parsed["pep8"], - "type_hints": parsed["type_hints"], - "edge_cases": parsed["edge_cases"], - "naming": parsed["naming"], + "pep8": result["pep8"], + "type_hints": result["type_hints"], + "edge_cases": result["edge_cases"], + "naming": result["naming"], }, - "weakest_criterion": parsed["weakest_criterion"], - "verdict": parsed["verdict"], + "weakest_criterion": result["weakest_criterion"], + "verdict": result["verdict"], } -# ---------- 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 +# Node: rewrite + +def rewrite(state: CodeReviewState) -> Dict[str, str]: + # Increment round + state["round"] += 1 + prompt = ChatPromptTemplate.from_messages([ + ("system", "You are a senior Python developer. Rewrite the review to improve the section about {weakest_criterion}. Keep other points unchanged."), + ("user", "Original review:\n{draft_review}") + ]) + chain = prompt | llm + new_review = chain.invoke({"weakest_criterion": state["weakest_criterion"], "draft_review": state["draft_review"]}) return {"draft_review": new_review} -# ---------- Graph ---------- +# Build the 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") - builder.add_edge("draft_review", "reflect") +# Conditional edge after reflect builder.add_conditional_edges( "reflect", - lambda x: "END" if x["verdict"] == "ok" else "rewrite", + lambda state: "END" if state["verdict"] == "ok" else "rewrite", ) builder.add_edge("rewrite", "reflect") -# Stop after max_rounds -builder.add_conditional_edges( - "reflect", - lambda x: "END" if x["round"] >= x["max_rounds"] else "rewrite", -) +builder.set_entry_point("draft_review") +builder.set_finish_point("END") graph = builder.compile() -# ---------- Demo ---------- -async def main(): - # Example function to review - code = """ +# Demo +if __name__ == "__main__": def sort_numbers(arr): return sorted(arr) - """ - init_state: CodeReviewState = { + + code = inspect.getsource(sort_numbers) + initial_state: CodeReviewState = { "code": code, "draft_review": "", "criteria_scores": {}, @@ -144,14 +99,12 @@ async def main(): "round": 0, "max_rounds": 2, } - result = await graph.ainvoke(init_state) + result = graph.invoke(initial_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()) + print("\n--- Round ---") + print(result["round"])