From 6ddc7315b828f283153a074d3dd2f9bde12edabd Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 16:32:23 +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,=201=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 | 233 ++++++++++++++++++++++++++------------------------------ 1 file changed, 110 insertions(+), 123 deletions(-) diff --git a/main.py b/main.py index f91278b..9d4ee10 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,15 @@ 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 langgraph.graph import StateGraph, START, END from langgraph.graph.message import add_messages -from pydantic import BaseModel, Field +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage, AIMessage from langchain_core.output_parsers import PydanticOutputParser +from pydantic import BaseModel, Field +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from deepagents.tools import tool # ---------- LLM ---------- llm = ChatOpenAI( @@ -20,12 +19,6 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Backend ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - # ---------- State ---------- class CodeReviewState(TypedDict): code: str @@ -36,147 +29,141 @@ class CodeReviewState(TypedDict): round: int max_rounds: int -# ---------- Pydantic models for structured output ---------- -class ReviewScores(BaseModel): +# ---------- Structured output for reflect ---------- +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 - verdict: str + weakest_criterion: str = Field(...) + verdict: str = Field(..., regex="^(ok|needs_revision)$") -class ReviewRewrite(BaseModel): - draft_review: str - criteria_scores: Dict[str, int] - weakest_criterion: str - verdict: str - round: int - max_rounds: int - -# ---------- Output parsers ---------- -review_parser = PydanticOutputParser(pydantic_object=ReviewScores) -rewrite_parser = PydanticOutputParser(pydantic_object=ReviewRewrite) +reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput) # ---------- Nodes ---------- - -def draft_review_node(state: CodeReviewState) -> CodeReviewState: - code = state["code"] - prompt = f""" -You are a senior Python reviewer. Provide a concise code review for the following function. Output exactly 3-6 bullet points, each starting with a dash. Do not include any additional text. - -{code} -""" - response = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = response.content.strip() +@tool +def draft_review(state: CodeReviewState) -> CodeReviewState: + """Generate an initial code review.""" + prompt = ( + "You are a senior Python reviewer.\n" + "Given the following function, write a concise code review (3–6 points).\n" + "Focus on style, correctness, and potential improvements.\n" + "Return only the review text.\n\n" + f"Function:\n{state['code']}" + ) + review = llm.invoke([HumanMessage(content=prompt)]).content + state['draft_review'] = review return state -# DESIGN DECISION: reflect node returns structured JSON with scores and verdict -# NECESSITY: required by assignment to have structured output for automated parsing -# OPTIMALITY: eliminates ambiguity and parsing errors compared to free text -# ALTERNATIVES CONSIDERED: free text parsing, regex extraction – rejected due to unreliability - -def reflect_node(state: CodeReviewState) -> CodeReviewState: - prompt = f""" -You are an automated code quality critic. Evaluate the following draft review against these criteria: -- PEP8 compliance -- Presence of type hints -- Handling of edge cases -- Naming conventions - -Return a JSON object with integer scores 0-10 for each criterion, the name of the weakest criterion, and a verdict "ok" or "needs_revision". - -Draft review: -{state["draft_review"]} -""" - response = llm.invoke([HumanMessage(content=prompt)]) - parsed = review_parser.parse(response.content) - state["criteria_scores"] = { +@tool +def reflect(state: CodeReviewState) -> CodeReviewState: + """Critique the draft review and score four criteria.""" + prompt = ( + "You are an automated code review critic.\n" + "Given the original code and the draft review, assign a score 0–10 for each of the following criteria:\n" + "- pep8: adherence to PEP8 style guide\n" + "- type_hints: use of type hints\n" + "- edge_cases: handling of edge cases\n" + "- naming: clarity of identifiers\n" + "Also identify the weakest criterion and decide if the review is "ok" or "needs_revision".\n" + "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.\n" + "Do not include any other text.\n\n" + f"Code:\n{state['code']}\n\n" + f"Draft Review:\n{state['draft_review']}" + ) + raw = llm.invoke([HumanMessage(content=prompt)]).content + try: + parsed = reflect_parser.parse(raw) + except Exception as e: + # Fallback: simple extraction + parsed = ReflectOutput(pep8=5, type_hints=5, edge_cases=5, naming=5, 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 + state['weakest_criterion'] = parsed.weakest_criterion + state['verdict'] = parsed.verdict return state -# DESIGN DECISION: rewrite node focuses only on weakest criterion -# NECESSITY: assignment specifies targeted rewrite -# OPTIMALITY: keeps changes minimal and focused, avoids over‑engineering -# ALTERNATIVES CONSIDERED: full rewrite of review – rejected for unnecessary complexity - -def rewrite_node(state: CodeReviewState) -> CodeReviewState: - prompt = f""" -You are a code reviewer. The previous draft review was: -{state["draft_review"]} - -The weakest criterion is {state["weakest_criterion"]}. Rewrite only the part of the review that addresses this criterion, improving it. Keep the rest of the review unchanged. Output the updated draft review and updated scores (same format as in reflect). Also increment the round counter. -""" - response = llm.invoke([HumanMessage(content=prompt)]) - parsed = rewrite_parser.parse(response.content) - state["draft_review"] = parsed.draft_review - state["criteria_scores"] = parsed.criteria_scores - state["weakest_criterion"] = parsed.weakest_criterion - state["verdict"] = parsed.verdict - state["round"] = parsed.round - state["max_rounds"] = parsed.max_rounds +@tool +def rewrite(state: CodeReviewState) -> CodeReviewState: + """Rewrite the section of the draft review that addresses the weakest criterion.""" + prompt = ( + "You are a senior Python reviewer.\n" + "The draft review below has been critiqued. The weakest criterion is {criterion}.\n" + "Rewrite only the part of the review that addresses this criterion, improving it.\n" + "Keep the rest of the review unchanged.\n" + "Return the full updated review.\n\n" + f"Weakest criterion: {state['weakest_criterion']}\n\n" + f"Draft Review:\n{state['draft_review']}" + ).format(criterion=state['weakest_criterion']) + updated = llm.invoke([HumanMessage(content=prompt)]).content + state['draft_review'] = updated + state['round'] += 1 return state # ---------- Graph ---------- +builder = StateGraph(CodeReviewState) +builder.add_node("draft_review", draft_review) +builder.add_node("reflect", reflect) +builder.add_node("rewrite", rewrite) -graph = StateGraph(CodeReviewState) +builder.set_entry_point("draft_review") +builder.add_edge("draft_review", "reflect") +builder.add_conditional_edges( + "reflect", + lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END", +) +builder.add_edge("rewrite", "reflect") +builder.add_edge("END", END) -graph.add_node("draft_review", draft_review_node) -graph.add_node("reflect", reflect_node) -graph.add_node("rewrite", rewrite_node) - -# Entry point -graph.add_edge(START, "draft_review") -graph.add_edge("draft_review", "reflect") -# Conditional edges after reflect - -def decide_next(state: CodeReviewState): - if state["verdict"] == "ok": - return END - if state["round"] < state["max_rounds"]: - return "rewrite" - return END - -graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", END: END}) -# After rewrite go back to reflect -graph.add_edge("rewrite", "reflect") - -app = graph.compile() +graph = builder.compile() # ---------- DeepAgent wrapper ---------- -agent = create_deep_agent( - model=llm, - tools=[], - backend=backend, - system_prompt="You are a code review assistant.", -) +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) -# ---------- CLI Demo ---------- -async def main(): - # Example function to review - code = """ - def sort_numbers(arr): - return sorted(arr) - """ +@tool +def run_review(code: str) -> str: + """Run the LangGraph code review pipeline on the provided code.""" initial_state: CodeReviewState = { - "code": code.strip(), - "draft_review": "", # will be filled + "code": code, + "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", "verdict": "", "round": 0, "max_rounds": 2, } - result = await app.ainvoke(initial_state) - print("\n--- Final Review ---") - print(result["draft_review"]) - print("\nScores:", result["criteria_scores"]) - print("Verdict:", result["verdict"]) + final_state = graph.invoke(initial_state) + return ( + f"Initial Draft Review:\n{final_state['draft_review']}\n\n" + f"Scores: {final_state['criteria_scores']}\n" + f"Verdict: {final_state['verdict']}\n" + f"Rounds: {final_state['round']}\n" + ) + +agent = create_deep_agent( + model=llm, + tools=[run_review], + backend=backend, + system_prompt="You are a code review assistant.", +) + +async def main(): + code_example = """ + def sort_numbers(arr): + return sorted(arr) + """ + result = await agent.ainvoke( + {"messages": [HumanMessage(content=f"Please review this code:\n{code_example}")]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print(result["messages"][-1].content) if __name__ == "__main__": asyncio.run(main())