From bf2de97ff888b7fbd7d46ce1e4f151b1dc65e766 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: Tue, 30 Jun 2026 18:50:49 +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 | 199 ++++++++++++++++++++++++++++++++------------------------ 1 file changed, 113 insertions(+), 86 deletions(-) diff --git a/main.py b/main.py index c2c2d43..35e425b 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,37 @@ +""" +# 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 langchain_openai import ChatOpenAI -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 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 pydantic import BaseModel, Field -from langchain_core.output_parsers import PydanticOutputParser +from dotenv import load_dotenv -# ---------- LLM ---------- +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 = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -21,51 +39,52 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- State ---------- -class CodeReviewState(TypedDict): - code: str - draft_review: str - criteria_scores: Dict[str, int] - weakest_criterion: str - verdict: str - round: int - max_rounds: int - -# ---------- Pydantic for reflect output ---------- +# --------------------------------------------------------------------------- +# Structured output models for reflect node +# --------------------------------------------------------------------------- class ReflectOutput(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 edge case handling") - naming: int = Field(..., description="Score 0-10 for naming conventions") - weakest_criterion: str = Field(..., description="Name of the weakest criterion") - verdict: str = Field(..., description="'ok' or 'needs_revision'") + 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)$") reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput) -# ---------- Nodes ---------- -async def draft_review(state: CodeReviewState) -> CodeReviewState: - prompt = f"""Please write a concise code review (3-6 bullet points) for the following Python function. Focus on style, type hints, edge cases, and naming. +# --------------------------------------------------------------------------- +# 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. -```python +Function code: {state['code']} -``` -Return only the review text.""" - review = await llm.ainvoke([HumanMessage(content=prompt)]) - state['draft_review'] = review.content.strip() +Review:""" + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = response.content.strip() return state -async def reflect(state: CodeReviewState) -> CodeReviewState: - prompt = f"""You are a code review critic. Evaluate the following review text and assign scores 0-10 for each of the four criteria: pep8, type_hints, edge_cases, naming. Also identify the weakest criterion and decide if the review is "ok" or "needs_revision". +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: -Review text: +{reflect_parser.get_format_instructions()} + +Draft review: {state['draft_review']} -Provide the output in the following JSON-like format: -{{"pep8": int, "type_hints": int, "edge_cases": int, "naming": int, "weakest_criterion": str, "verdict": str}} -""" - raw = await llm.ainvoke([HumanMessage(content=prompt)]) - parsed = reflect_parser.parse(raw.content) +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, @@ -76,35 +95,47 @@ Provide the output in the following JSON-like format: state['verdict'] = parsed.verdict return state -async def rewrite(state: CodeReviewState) -> CodeReviewState: - # Simple rewrite: add a sentence addressing the weakest criterion - additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}.") - state['draft_review'] = state['draft_review'] + "\n" + additional +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 ---------- -def build_graph() -> StateGraph[CodeReviewState]: +# --------------------------------------------------------------------------- +# Graph construction +# --------------------------------------------------------------------------- +def create_graph() -> StateGraph: graph = StateGraph(CodeReviewState) - graph.add_node("draft_review", draft_review) - graph.add_node("reflect", reflect) - graph.add_node("rewrite", rewrite) + 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 x: "END" if x['verdict'] == "ok" or x['round'] >= x['max_rounds'] else "rewrite", + lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END", ) graph.add_edge("rewrite", "reflect") - return graph.compile() + return graph -# ---------- Tool ---------- -@tool -def code_review_tool(code: str) -> str: - """Perform a structured code review with possible rewrites.""" - graph = build_graph() +# --------------------------------------------------------------------------- +# CLI helper +# --------------------------------------------------------------------------- +async def run_review(code: str, max_rounds: int = 2) -> CodeReviewState: initial_state: CodeReviewState = { "code": code, "draft_review": "", @@ -112,34 +143,30 @@ def code_review_tool(code: str) -> str: "weakest_criterion": "", "verdict": "", "round": 0, - "max_rounds": 2, + "max_rounds": max_rounds, } - final_state = graph.invoke(initial_state) - return f"Final Review:\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}" - -# ---------- DeepAgent ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -agent = create_deep_agent( - model=llm, - tools=[code_review_tool], - backend=backend, - system_prompt="You are a helpful code review assistant.", -) - -async def main(): - sample_code = """ - def sort_numbers(arr): - return sorted(arr) - """ - result = await agent.ainvoke( - {"messages": [HumanMessage(content=f"Please review this function:\n{sample_code}")]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print(result["messages"][-1].content) + graph = create_graph() + final_state = await graph.astate(initial_state) + return final_state +# --------------------------------------------------------------------------- +# Demo main +# --------------------------------------------------------------------------- if __name__ == "__main__": - asyncio.run(main()) + sample_code = """ +def sort_numbers(arr): + return sorted(arr) +""" + result = asyncio.run(run_review(sample_code)) + print("\n=== Initial Draft 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