From ca177e5e940542ac012814cde627442b6af5f49b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 09:16:25 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=9F=D0=BE=D0=B2?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BD=D1=8B=D0=B9=20=D1=8D=D0=BA=D0=B7=D0=B0?= =?UTF-8?q?=D0=BC=D0=B5=D0=BD=20#2:=20=D0=93=D1=80=D0=B0=D1=84=20=D1=81=20?= =?UTF-8?q?=D1=80=D0=B5=D1=84=D0=BB=D0=B5=D0=BA=D1=81=D0=B8=D0=B5=D0=B9=20?= =?UTF-8?q?=D0=BD=D0=B0=20=D0=BA=D0=BE=D0=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 203 +++++++++++++++++++++++++++----------------------------- 1 file changed, 98 insertions(+), 105 deletions(-) diff --git a/main.py b/main.py index 107f02e..478877d 100644 --- a/main.py +++ b/main.py @@ -1,38 +1,18 @@ import os import asyncio -from typing import TypedDict, Dict +from typing import TypedDict, Annotated, Dict -from dotenv import load_dotenv 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_core.output_parsers import PydanticOutputParser -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend - -# Load API key from .env -load_dotenv() - -# LLM configuration - OpenRouter -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=0.0, -) - -# Backend for deepagents - simple filesystem -backend = FilesystemBackend() - -# Create a deepagents agent that will be used inside the graph nodes -agent = create_deep_agent( - model=llm, - tools=[], - backend=backend, - system_prompt="You are a code review assistant.", -) # ---------- State definition ---------- class CodeReviewState(TypedDict): @@ -46,117 +26,127 @@ class CodeReviewState(TypedDict): # ---------- Structured output for critic ---------- class CriticOutput(BaseModel): - scores: Dict[str, int] = Field( - description="Scores for each criterion: pep8, type_hints, edge_cases, naming. Values 0-10." - ) - verdict: str = Field( - description='Verdict: "ok" if all scores >= 7, otherwise "needs_revision".' - ) + pep8: int = Field(description="Score for PEP8 compliance (0-10)") + type_hints: int = Field(description="Score for type hints (0-10)") + edge_cases: int = Field(description="Score for edge case handling (0-10)") + naming: int = Field(description="Score for naming conventions (0-10)") + verdict: str = Field(description='Verdict: "ok" or "needs_revision"') critic_parser = PydanticOutputParser(pydantic_object=CriticOutput) -# ---------- Graph nodes ---------- +# ---------- LLM and backend ---------- +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, +) + +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) + +agent = create_deep_agent( + model=llm, + tools=[], + backend=backend, + system_prompt="You are a helpful agent.", +) + +# ---------- Node functions ---------- async def draft_review(state: CodeReviewState) -> CodeReviewState: - prompt = ( - f"Write a concise code review (3-6 bullet points) for the following Python function:\n\n" - f"{state['code']}\n\n" - "Focus on style, correctness, and potential improvements." - ) - response = await agent.ainvoke( - {"messages": [HumanMessage(content=prompt)]}, - {"configurable": {"thread_id": "draft_review"}}, - ) - review_text = response["messages"][-1].content.strip() - state["draft_review"] = review_text + prompt = f"Please provide a concise code review (3-6 points) for the following Python function:\n\n{state['code']}" + result = await agent.ainvoke([HumanMessage(content=prompt)]) + review = result["messages"][-1].content.strip() + state["draft_review"] = review print("\n--- Draft Review ---") - print(review_text) + print(review) return state async def reflect(state: CodeReviewState) -> CodeReviewState: prompt = ( - f"Evaluate the following code review and assign scores (0-10) for each criterion:\n\n" - f"Review:\n{state['draft_review']}\n\n" - "Criteria:\n" - "1. pep8: adherence to PEP8 style guide.\n" - "2. type_hints: presence and correctness of type hints.\n" - "3. edge_cases: handling of edge cases and robustness.\n" - "4. naming: clarity and consistency of names.\n\n" - "Return a JSON object with keys 'scores' (dict) and 'verdict' ('ok' or 'needs_revision')." + f"Evaluate the following draft review:\n\n{state['draft_review']}\n\n" + "Score each of the following criteria on a scale of 0-10:\n" + "- pep8\n- type_hints\n- edge_cases\n- naming\n\n" + "Return the scores and a verdict ('ok' or 'needs_revision') in the following JSON format:\n" + "{\n \"pep8\": int,\n \"type_hints\": int,\n \"edge_cases\": int,\n \"naming\": int,\n \"verdict\": \"ok\" | \"needs_revision\"\n}" ) - response = await agent.ainvoke( - {"messages": [HumanMessage(content=prompt)]}, - {"configurable": {"thread_id": "reflect"}}, - ) - raw_output = response["messages"][-1].content.strip() + result = await agent.ainvoke([HumanMessage(content=prompt)]) + raw_output = result["messages"][-1].content.strip() try: parsed = critic_parser.parse(raw_output) except Exception as e: - # Fallback: simple parsing if JSON is malformed - import json - parsed = CriticOutput(**json.loads(raw_output)) - state["criteria_scores"] = parsed.scores - # Determine weakest criterion - weakest = min(parsed.scores.items(), key=lambda kv: kv[1])[0] + # Fallback: simple parsing if LLM output is not perfectly formatted + parsed = CriticOutput( + pep8=0, + type_hints=0, + edge_cases=0, + naming=0, + verdict="needs_revision", + ) + scores = { + "pep8": parsed.pep8, + "type_hints": parsed.type_hints, + "edge_cases": parsed.edge_cases, + "naming": parsed.naming, + } + weakest = min(scores, key=scores.get) + state["criteria_scores"] = scores state["weakest_criterion"] = weakest state["verdict"] = parsed.verdict - print("\n--- Critic Scores ---") - for crit, score in parsed.scores.items(): - print(f"{crit}: {score}") + print("\n--- Reflection ---") + print(f"Scores: {scores}") print(f"Weakest criterion: {weakest}") print(f"Verdict: {parsed.verdict}") return state async def rewrite(state: CodeReviewState) -> CodeReviewState: - state["round"] += 1 prompt = ( - f"Rewrite the part of the review that addresses the weakest criterion " - f"('{state['weakest_criterion']}') to improve it. Keep the rest of the review unchanged.\n\n" - f"Original Review:\n{state['draft_review']}\n\n" - "Provide only the updated review." + f"Rewrite the section of the draft review that addresses the weakest criterion " + f"('{state['weakest_criterion']}') to improve it. Keep all other parts unchanged.\n\n" + f"Original draft review:\n\n{state['draft_review']}" ) - response = await agent.ainvoke( - {"messages": [HumanMessage(content=prompt)]}, - {"configurable": {"thread_id": "rewrite"}}, - ) - new_review = response["messages"][-1].content.strip() + result = await agent.ainvoke([HumanMessage(content=prompt)]) + new_review = result["messages"][-1].content.strip() state["draft_review"] = new_review - print("\n--- Rewritten Review (Round {}) ---".format(state["round"])) + state["round"] += 1 + print("\n--- Rewritten Review ---") print(new_review) return state -# ---------- Graph construction ---------- -builder = StateGraph(CodeReviewState) +# ---------- Graph ---------- +def build_graph() -> StateGraph: + graph = StateGraph(CodeReviewState) + graph.add_node("draft_review", draft_review) + graph.add_node("reflect", reflect) + graph.add_node("rewrite", rewrite) -builder.add_node("draft_review", draft_review) -builder.add_node("reflect", reflect) -builder.add_node("rewrite", rewrite) + graph.add_edge(START, "draft_review") + graph.add_edge("draft_review", "reflect") -builder.add_edge(START, "draft_review") -builder.add_edge("draft_review", "reflect") - -# Conditional edges after reflect -def reflect_conditional(state: CodeReviewState): - if state["verdict"] == "ok": + def reflect_cond(state: CodeReviewState): + if state["verdict"] == "ok": + return END + if state["round"] < state["max_rounds"]: + return "rewrite" return END - if state["round"] < state["max_rounds"]: - return "rewrite" - return END -builder.add_conditional_edges("reflect", reflect_conditional) + graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", END: END}) + graph.add_edge("rewrite", "reflect") -builder.add_edge("rewrite", "reflect") - -graph = builder.compile() + return graph # ---------- Demo ---------- async def main(): # Sample function to review - sample_code = """ -def sort_numbers(arr): - return sorted(arr) -""" + code_str = """def sort_numbers(arr): + return sorted(arr)""" + initial_state: CodeReviewState = { - "code": sample_code.strip(), + "code": code_str, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -165,12 +155,15 @@ def sort_numbers(arr): "max_rounds": 2, } - final_state = await graph.ainvoke(initial_state) + graph = build_graph() + app = graph.compile() + final_state = await app.ainvoke(initial_state) + print("\n=== Final State ===") + print(f"Round: {final_state['round']}") print(f"Verdict: {final_state['verdict']}") - print(f"Rounds performed: {final_state['round']}") - print("\nFinal Review:") - print(final_state["draft_review"]) + print(f"Draft Review:\n{final_state['draft_review']}") + print(f"Scores: {final_state['criteria_scores']}") if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file