From 25cc4c772d5765d5fc9f0c31742ab83f43f0b3c9 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 13:23:17 +0000 Subject: [PATCH] Rebuild the task as a single LangGraph code-review reflection graph with structured scores, targeted rewrite, max_rounds guard, and CLI demo.: update main.py --- main.py | 252 ++++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 173 insertions(+), 79 deletions(-) diff --git a/main.py b/main.py index b0720e4..4f6147b 100644 --- a/main.py +++ b/main.py @@ -1,97 +1,168 @@ -from typing import TypedDict, Dict -import inspect +import os +import textwrap +from typing import Literal, TypedDict + +from dotenv import load_dotenv +from langchain_openai import ChatOpenAI +from langgraph.graph import END, START, StateGraph +from pydantic import BaseModel, Field -from langgraph.graph import StateGraph, END -from langchain_ollama import Ollama -from langchain_core.prompts import ChatPromptTemplate -from langchain_core.output_parsers import JsonOutputParser -# Define the state class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: Dict[str, int] + criteria_scores: dict[str, int] weakest_criterion: str - verdict: str # "ok" | "needs_revision" + verdict: str round: int max_rounds: int -# LLM instance (Ollama) -llm = Ollama(model="llama3.1") -# Node: draft_review +class ReflectionResult(BaseModel): + criteria_scores: dict[str, int] = Field( + description="Scores from 0 to 10 for pep8, type_hints, edge_cases, naming." + ) + weakest_criterion: Literal["pep8", "type_hints", "edge_cases", "naming"] + verdict: Literal["ok", "needs_revision"] -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} -# Node: reflect +def build_llm() -> ChatOpenAI: + return ChatOpenAI( + model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b"), + base_url=os.getenv("OPENAI_BASE_URL"), + api_key=os.getenv("OPENAI_API_KEY", "dummy"), + temperature=0, + ) -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 + +def draft_review(state: CodeReviewState) -> dict: + llm = build_llm() + prompt = textwrap.dedent( + """ + Ты опытный Python code reviewer. + Напиши code review по функции ниже. + Требования: + - 3-6 конкретных пунктов; + - оцени сильные стороны и что улучшить; + - обязательно затронь PEP8, type hints, edge cases и naming, если это уместно; + - ответ пиши на русском. + + Код: + {code} + """ + ).strip() + response = llm.invoke(prompt.format(code=state["code"])) + return {"draft_review": response.content.strip()} + + +def reflect(state: CodeReviewState) -> dict: + llm = build_llm().with_structured_output(ReflectionResult) + prompt = textwrap.dedent( + """ + Ты критик качества code review. + Оцени review по 4 критериям: + - pep8 + - type_hints + - edge_cases + - naming + + Правила: + - для каждого критерия выставь integer score от 0 до 10; + - weakest_criterion — самый слабый критерий; + - verdict = "ok", если review уже достаточно хорошее; + - verdict = "needs_revision", если самое слабое место стоит усилить. + + Код: + {code} + + Review: + {review} + """ + ).strip() + result = llm.invoke( + prompt.format(code=state["code"], review=state["draft_review"]) + ) return { - "criteria_scores": { - "pep8": result["pep8"], - "type_hints": result["type_hints"], - "edge_cases": result["edge_cases"], - "naming": result["naming"], - }, - "weakest_criterion": result["weakest_criterion"], - "verdict": result["verdict"], + "criteria_scores": result.criteria_scores, + "weakest_criterion": result.weakest_criterion, + "verdict": result.verdict, } -# 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} +def rewrite(state: CodeReviewState) -> dict: + llm = build_llm() + prompt = textwrap.dedent( + """ + Ты улучшаешь уже написанный code review. + Нужно целенаправленно усилить самое слабое место: {weakest_criterion}. -# Build the graph -builder = StateGraph(CodeReviewState) -builder.add_node("draft_review", draft_review) -builder.add_node("reflect", reflect) -builder.add_node("rewrite", rewrite) + Требования: + - сохрани общий формат краткого review; + - сделай акцент именно на критерии {weakest_criterion}; + - добавь более точные и полезные замечания; + - ответ пиши на русском; + - итог должен остаться в формате 3-6 пунктов. -builder.add_edge("draft_review", "reflect") -# Conditional edge after reflect -builder.add_conditional_edges( - "reflect", - lambda state: "END" if state["verdict"] == "ok" else "rewrite", -) -builder.add_edge("rewrite", "reflect") + Код: + {code} -builder.set_entry_point("draft_review") -builder.set_finish_point("END") + Текущий review: + {review} + """ + ).strip() + response = llm.invoke( + prompt.format( + weakest_criterion=state["weakest_criterion"], + code=state["code"], + review=state["draft_review"], + ) + ) + return { + "draft_review": response.content.strip(), + "round": state["round"] + 1, + } -graph = builder.compile() -# Demo -if __name__ == "__main__": - def sort_numbers(arr): - return sorted(arr) +def next_step(state: CodeReviewState) -> str: + if state["verdict"] == "ok": + return "finish" + if state["round"] >= state["max_rounds"]: + return "finish" + return "rewrite" - code = inspect.getsource(sort_numbers) + +def build_graph(): + builder = StateGraph(CodeReviewState) + builder.add_node("draft_review", draft_review) + builder.add_node("reflect", reflect) + builder.add_node("rewrite", rewrite) + builder.add_edge(START, "draft_review") + builder.add_edge("draft_review", "reflect") + builder.add_conditional_edges( + "reflect", + next_step, + { + "rewrite": "rewrite", + "finish": END, + }, + ) + builder.add_edge("rewrite", "reflect") + return builder.compile() + + +def demo_code() -> str: + return textwrap.dedent( + """ + def sort_numbers(arr): + return sorted(arr) + """ + ).strip() + + +def run_demo() -> None: + graph = build_graph() initial_state: CodeReviewState = { - "code": code, + "code": demo_code(), "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -99,12 +170,35 @@ if __name__ == "__main__": "round": 0, "max_rounds": 2, } - 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"]) - print("\n--- Round ---") - print(result["round"]) + + print("=== Code Review Reflection Demo ===") + print("Code under review:") + print(initial_state["code"]) + + final_state = initial_state.copy() + for chunk in graph.stream(initial_state, stream_mode="updates"): + for node_name, update in chunk.items(): + final_state.update(update) + print() + if node_name == "draft_review": + print("--- Draft Review ---") + print(final_state["draft_review"]) + elif node_name == "reflect": + print("--- Critic Scores ---") + print(final_state["criteria_scores"]) + print(f"Weakest criterion: {final_state['weakest_criterion']}") + print(f"Verdict: {final_state['verdict']}") + elif node_name == "rewrite": + print(f"--- Rewritten Review After Round {final_state['round']} ---") + print(final_state["draft_review"]) + + print() + print("=== Final Result ===") + print(final_state["draft_review"]) + print(final_state["criteria_scores"]) + print(f"Rounds used: {final_state['round']} / {final_state['max_rounds']}") + + +if __name__ == "__main__": + load_dotenv() + run_demo()