commit ddb8e621b03722a526eac1cd84e7a155e2dc31e7 Author: Илья 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Sat Jun 27 13:38:39 2026 +0000 add: main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..c2c2d43 --- /dev/null +++ b/main.py @@ -0,0 +1,145 @@ +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, +) + +# ---------- 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 ---------- +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'") + +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. + +```python +{state['code']} +``` + +Return only the review text.""" + review = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = review.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". + +Review text: +{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) + 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: + # 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 + state['round'] += 1 + return state + +# ---------- Graph ---------- +def build_graph() -> StateGraph[CodeReviewState]: + graph = StateGraph(CodeReviewState) + graph.add_node("draft_review", draft_review) + graph.add_node("reflect", reflect) + graph.add_node("rewrite", rewrite) + + graph.set_entry_point("draft_review") + 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", + ) + graph.add_edge("rewrite", "reflect") + + return graph.compile() + +# ---------- Tool ---------- +@tool +def code_review_tool(code: str) -> str: + """Perform a structured code review with possible rewrites.""" + graph = build_graph() + initial_state: CodeReviewState = { + "code": code, + "draft_review": "", + "criteria_scores": {}, + "weakest_criterion": "", + "verdict": "", + "round": 0, + "max_rounds": 2, + } + 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) + +if __name__ == "__main__": + asyncio.run(main())