commit 1cafdb65166b04bb4f44a96522390fc61e635cf4 Author: Даниил Викторов Date: Tue Jun 30 05:49:21 2026 +0000 add: main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..eabdb26 --- /dev/null +++ b/main.py @@ -0,0 +1,183 @@ +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_core.output_parsers import PydanticOutputParser + +# ---------- LLM ---------- +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 ---------- +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# ---------- 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 model for reflect output ---------- +class ReflectOutput(BaseModel): + 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)") + weakest_criterion: str = Field(description="The criterion with the lowest score") + verdict: str = Field(description="'ok' or 'needs_revision'") + +parser = PydanticOutputParser(pydantic_object=ReflectOutput) + +# ---------- Nodes ---------- +async def draft_review(state: CodeReviewState) -> CodeReviewState: + prompt = f""" + Write a concise code review for the following Python function. Provide 3-6 bullet points. + + Function: + ```python + {state['code']} + ``` + """ + response = await llm.ainvoke(HumanMessage(content=prompt)) + state['draft_review'] = response.content.strip() + return state + +async def reflect(state: CodeReviewState) -> CodeReviewState: + prompt = f""" + Evaluate the following code and draft review. Score each of the four criteria on a scale 0-10: + - PEP8 compliance + - Type hints + - Edge case handling + - Naming conventions + + Provide the scores, identify the weakest criterion, and give a verdict ('ok' or 'needs_revision'). + + Code: + ```python + {state['code']} + ``` + + Draft Review: + ```text + {state['draft_review']} + ``` + + Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict. + """ + response = await llm.ainvoke(HumanMessage(content=prompt)) + try: + parsed = parser.parse(response.content) + except Exception as e: + # Fallback: set all scores to 0 and verdict to needs_revision + parsed = ReflectOutput(pep8=0, type_hints=0, edge_cases=0, naming=0, weakest_criterion="unknown", 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(state: CodeReviewState) -> CodeReviewState: + prompt = f""" + Rewrite the section of the draft review that addresses the weakest criterion: {state['weakest_criterion']}. + Keep all other parts of the review unchanged. + + Current Draft Review: + ```text + {state['draft_review']} + ``` + """ + response = await llm.ainvoke(HumanMessage(content=prompt)) + state['draft_review'] = response.content.strip() + state['round'] += 1 + return state + +# ---------- Graph ---------- +review_graph = StateGraph(CodeReviewState) +review_graph.add_node("draft_review", draft_review) +review_graph.add_node("reflect", reflect) +review_graph.add_node("rewrite", rewrite) + +review_graph.set_entry_point("draft_review") +review_graph.add_edge("draft_review", "reflect") + +# Conditional edges after reflect +review_graph.add_conditional_edges( + "reflect", + lambda state: "END" if state["verdict"] == "ok" else "rewrite" if state["round"] < state["max_rounds"] else "END", +) +review_graph.add_edge("rewrite", "reflect") + +compiled_graph = review_graph.compile() + +# ---------- Tool ---------- +@tool +async def run_code_review(code: str) -> Dict: + """Run a code review on the provided Python function.""" + initial_state: CodeReviewState = { + "code": code, + "draft_review": "", + "criteria_scores": {}, + "weakest_criterion": "", + "verdict": "", + "round": 0, + "max_rounds": 2, + } + final_state = compiled_graph.invoke(initial_state) + return { + "final_review": final_state["draft_review"], + "scores": final_state["criteria_scores"], + "weakest_criterion": final_state["weakest_criterion"], + "verdict": final_state["verdict"], + "rounds": final_state["round"], + } + +# ---------- Agent ---------- +agent = create_deep_agent( + model=llm, + tools=[run_code_review], + backend=backend, + system_prompt="You are a helpful code review agent.", +) + +# ---------- Demo ---------- +async def main(): + # Sample function to review + code = """ + def sort_numbers(arr): + return sorted(arr) + """ + result = await run_code_review(code) + print("\n=== Final Review ===\n") + print(result["final_review"]) + print("\n=== Scores ===\n") + for k, v in result["scores"].items(): + print(f"{k}: {v}") + print(f"\nVerdict: {result['verdict']} (Rounds: {result['rounds']})") + +if __name__ == "__main__": + asyncio.run(main())