From 56646345a4c450c71b62a5e0909b3b8dce911231 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Wed, 1 Jul 2026 20:11:21 +0000 Subject: [PATCH] =?UTF-8?q?fix(needs=5Ffixes):=201=20=D0=B8=D1=81=D0=BF?= =?UTF-8?q?=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D0=B9,=200=20=D0=BE?= =?UTF-8?q?=D1=82=D1=81=D1=82=D0=BE=D1=8F=D0=BD=D0=BE=20=E2=80=94=20main.p?= =?UTF-8?q?y?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 256 +++++++++++++++++++++++++------------------------------- 1 file changed, 114 insertions(+), 142 deletions(-) diff --git a/main.py b/main.py index 35e425b..40b96a9 100644 --- a/main.py +++ b/main.py @@ -1,37 +1,15 @@ -""" -# main.py -# LangGraph code review agent with reflection and rewrite loop. -# Implements the task specification without any deepagents dependency. -# Uses OpenRouter via langchain-openai. - import os import asyncio -from typing import TypedDict, Annotated, Dict +from typing import TypedDict, Annotated + from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage, AIMessage from langchain_core.output_parsers import PydanticOutputParser from langchain_core.pydantic_v1 import BaseModel, Field from langgraph.graph import StateGraph, START, END from langgraph.graph.message import add_messages -from dotenv import load_dotenv -load_dotenv() - -# --------------------------------------------------------------------------- -# State definition -# --------------------------------------------------------------------------- -class CodeReviewState(TypedDict): - code: str - draft_review: str - criteria_scores: Dict[str, int] - weakest_criterion: str - verdict: str # "ok" | "needs_revision" - round: int - max_rounds: int - -# --------------------------------------------------------------------------- -# LLM setup -# --------------------------------------------------------------------------- +# ---------- LLM ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -39,134 +17,128 @@ llm = ChatOpenAI( temperature=0.0, ) -# --------------------------------------------------------------------------- -# Structured output models for reflect node -# --------------------------------------------------------------------------- -class ReflectOutput(BaseModel): - pep8: int = Field(..., ge=0, le=10) - type_hints: int = Field(..., ge=0, le=10) - edge_cases: int = Field(..., ge=0, le=10) - naming: int = Field(..., ge=0, le=10) - weakest_criterion: str = Field(...) - verdict: str = Field(..., regex="^(ok|needs_revision)$") +# ---------- State ---------- +class CodeReviewState(TypedDict): + code: str + draft_review: str + criteria_scores: dict[str, int] + weakest_criterion: str + verdict: str # "ok" | "needs_revision" + round: int + max_rounds: int -reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput) +# ---------- Structured output for reflect ---------- +class ReflectionOutput(BaseModel): + pep8: int = Field(description="Score 0-10 for PEP8 compliance") + type_hints: int = Field(description="Score 0-10 for type hints usage") + edge_cases: int = Field(description="Score 0-10 for handling edge cases") + naming: int = Field(description="Score 0-10 for naming conventions") + weakest_criterion: str = Field(description="Criterion with lowest score") + verdict: str = Field(description="'ok' or 'needs_revision'") -# --------------------------------------------------------------------------- -# Node implementations -# --------------------------------------------------------------------------- -async def draft_review_node(state: CodeReviewState) -> CodeReviewState: - """Generate an initial code review with 3–6 bullet points.""" - prompt = f""" -You are a senior Python developer. You will write a concise code review for the following function. Provide 3 to 6 bullet points, each starting with a dash. +parser = PydanticOutputParser(pydantic_object=ReflectionOutput) -Function code: -{state['code']} +# ---------- Nodes ---------- -Review:""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state['draft_review'] = response.content.strip() - return state - -async def reflect_node(state: CodeReviewState) -> CodeReviewState: - """Critic evaluates the draft review on 4 criteria and returns structured scores.""" - prompt = f""" -You are a code review critic. Evaluate the following draft review on the four criteria below, assigning a score from 0 (worst) to 10 (excellent). Return the scores and the weakest criterion in a JSON format matching the schema: - -{reflect_parser.get_format_instructions()} - -Draft review: -{state['draft_review']} - -Scores:""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - try: - parsed = reflect_parser.parse(response.content) - except Exception as e: - # Fallback: treat as all zeros - parsed = ReflectOutput(pep8=0, type_hints=0, edge_cases=0, naming=0, weakest_criterion="pep8", verdict="needs_revision") - state['criteria_scores'] = { - "pep8": parsed.pep8, - "type_hints": parsed.type_hints, - "edge_cases": parsed.edge_cases, - "naming": parsed.naming, - } - state['weakest_criterion'] = parsed.weakest_criterion - state['verdict'] = parsed.verdict - return state - -async def rewrite_node(state: CodeReviewState) -> CodeReviewState: - """Rewrite the part of the review that addresses the weakest criterion.""" - prompt = f""" -You are a senior Python developer. The following code review has been identified as weak in the criterion: {state['weakest_criterion']}. Rewrite only the section of the review that addresses this criterion, improving clarity and depth. Keep the rest of the review unchanged. - -Original review: -{state['draft_review']} - -Rewritten review:""" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - # Replace only the weak section. For simplicity, we replace the whole review. - state['draft_review'] = response.content.strip() - state['round'] += 1 - return state - -# --------------------------------------------------------------------------- -# Graph construction -# --------------------------------------------------------------------------- -def create_graph() -> StateGraph: - graph = StateGraph(CodeReviewState) - graph.add_node("draft_review", draft_review_node) - graph.add_node("reflect", reflect_node) - graph.add_node("rewrite", rewrite_node) - - # Entry point - graph.set_entry_point("draft_review") - - # 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", +def draft_review_node(state: CodeReviewState) -> CodeReviewState: + code = state["code"] + prompt = ( + "You are a senior Python developer.\n" + "Given the following function, write a concise code review (3-6 bullet points).\n" + "Focus on style, correctness, edge cases, and naming.\n" + f"Function:\n{code}\n\nReview:" # LLM will output review ) - graph.add_edge("rewrite", "reflect") + response = llm.invoke([HumanMessage(content=prompt)]) + state["draft_review"] = response.content + return state - return graph -# --------------------------------------------------------------------------- -# CLI helper -# --------------------------------------------------------------------------- -async def run_review(code: str, max_rounds: int = 2) -> CodeReviewState: +def reflect_node(state: CodeReviewState) -> CodeReviewState: + review = state["draft_review"] + code = state["code"] + prompt = ( + "You are an automated code review critic.\n" + "Given the code and its review, assign a score 0-10 for each of the following criteria:\n" + "- pep8: PEP8 compliance\n" + "- type_hints: use of type hints\n" + "- edge_cases: handling of edge cases\n" + "- naming: clarity of names\n" + "Return the scores, the weakest criterion, and a verdict ('ok' if all scores >=7, else 'needs_revision').\n" + f"Code:\n{code}\n\nReview:\n{review}\n\nOutput in JSON with fields: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict." + ) + response = llm.invoke([HumanMessage(content=prompt)]) + try: + out = parser.parse(response.content) + except Exception as e: + # Fallback: simple parsing if JSON is not strict + import json + out = json.loads(response.content) + state["criteria_scores"] = { + "pep8": out.pep8, + "type_hints": out.type_hints, + "edge_cases": out.edge_cases, + "naming": out.naming, + } + state["weakest_criterion"] = out.weakest_criterion + state["verdict"] = out.verdict + return state + + +def rewrite_node(state: CodeReviewState) -> CodeReviewState: + weakest = state["weakest_criterion"] + review = state["draft_review"] + code = state["code"] + prompt = ( + "You are a senior Python developer tasked with improving a code review.\n" + f"The current review is:\n{review}\n\nThe weakest criterion is '{weakest}'.\n" + "Rewrite only the part of the review that addresses this criterion, making it stronger and more specific.\n" + "Keep the rest of the review unchanged.\n" + "Output only the updated review." + ) + response = llm.invoke([HumanMessage(content=prompt)]) + state["draft_review"] = response.content + state["round"] = state.get("round", 0) + 1 + return state + +# ---------- Graph ---------- +builder = StateGraph(CodeReviewState) +builder.add_node("draft_review", draft_review_node) +builder.add_node("reflect", reflect_node) +builder.add_node("rewrite", rewrite_node) + +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") +builder.add_edge("END", END) + +graph = builder.compile() + +# ---------- Demo ---------- +async def main(): + demo_code = """ + def sort_numbers(arr): + return sorted(arr) + """ initial_state: CodeReviewState = { - "code": code, - "draft_review": "", + "code": demo_code.strip(), + "draft_review": "", # will be filled "criteria_scores": {}, "weakest_criterion": "", "verdict": "", "round": 0, - "max_rounds": max_rounds, + "max_rounds": 2, } - graph = create_graph() - final_state = await graph.astate(initial_state) - return final_state - -# --------------------------------------------------------------------------- -# Demo main -# --------------------------------------------------------------------------- -if __name__ == "__main__": - sample_code = """ -def sort_numbers(arr): - return sorted(arr) -""" - result = asyncio.run(run_review(sample_code)) - print("\n=== Initial Draft Review ===") + result = await graph.ainvoke(initial_state) + print("\n--- Final Review ---") print(result["draft_review"]) - print("\n=== Scores ===") - print(result["criteria_scores"]) - print("\n=== Verdict ===") - print(result["verdict"]) - if result["verdict"] == "needs_revision": - print("\n=== Rewritten Review ===") - print(result["draft_review"]) # after last rewrite - print("\n=== Updated Scores ===") - print(result["criteria_scores"]) -"" \ No newline at end of file + print("\n--- Scores ---") + for k, v in result["criteria_scores"].items(): + print(f"{k}: {v}") + print(f"Verdict: {result['verdict']}") + +if __name__ == "__main__": + asyncio.run(main())