diff --git a/main.py b/main.py index 9d4ee10..191d774 100644 --- a/main.py +++ b/main.py @@ -1,136 +1,84 @@ import os import asyncio -from typing import TypedDict, Annotated, Dict -from langgraph.graph import StateGraph, START, END -from langgraph.graph.message import add_messages from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, AIMessage -from langchain_core.output_parsers import PydanticOutputParser +from langchain_core.messages import HumanMessage, SystemMessage +from langgraph.graph import StateGraph, START, END +from typing import TypedDict, Annotated +from langgraph.graph.message import add_messages from pydantic import BaseModel, Field -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from deepagents.tools import tool +from langchain_core.output_parsers import PydanticOutputParser -# ---------- LLM ---------- +# LLM setup llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", + model="gpt-4o-mini", base_url="https://openrouter.ai/api/v1", api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0, ) -# ---------- State ---------- +# State definition class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: Dict[str, int] + criteria_scores: dict[str, int] weakest_criterion: str verdict: str round: int max_rounds: int -# ---------- Structured output for reflect ---------- +# Node: draft_review +async def draft_review(state: CodeReviewState): + prompt = f"Write a concise code review (3-6 points) for the following Python function:\n\n{state['code']}" + msg = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = msg.content + return state + +# Node: 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 = Field(...) - verdict: str = Field(..., regex="^(ok|needs_revision)$") + scores: dict[str, int] + weakest: str + verdict: str -reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput) +parser = PydanticOutputParser(pydantic_object=ReflectOutput) -# ---------- Nodes ---------- -@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 +async def reflect(state: CodeReviewState): + prompt = f"Evaluate the draft review and assign scores 0-10 for PEP8, type_hints, edge_cases, naming. Return JSON with keys scores, weakest, verdict (ok or needs_revision).\n\nDraft review:\n{state['draft_review']}" + msg = await llm.ainvoke([HumanMessage(content=prompt)]) + out = parser.parse(msg.content) + state['criteria_scores'] = out.scores + state['weakest_criterion'] = out.weakest + state['verdict'] = out.verdict return state -@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 - return state - -@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 +# Node: rewrite +async def rewrite(state: CodeReviewState): + prompt = f"Rewrite the part of the draft review that addresses the weakest criterion '{state['weakest_criterion']}'. Keep other points unchanged.\n\nOriginal draft:\n{state['draft_review']}" + msg = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = msg.content 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 +graph = StateGraph(CodeReviewState) +graph.add_node("draft_review", draft_review) +graph.add_node("reflect", reflect) +graph.add_node("rewrite", rewrite) -builder.set_entry_point("draft_review") -builder.add_edge("draft_review", "reflect") -builder.add_conditional_edges( +graph.set_entry_point("draft_review") +graph.add_edge("draft_review", "reflect") +graph.add_conditional_edges( "reflect", - lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END", + lambda s: "rewrite" if s['verdict']=='needs_revision' and s['round'] str: - """Run the LangGraph code review pipeline on the provided code.""" - initial_state: CodeReviewState = { +# Demo function +async def demo(): + code = """def sort_numbers(arr): + return sorted(arr)""" + state: CodeReviewState = { "code": code, "draft_review": "", "criteria_scores": {}, @@ -139,31 +87,10 @@ def run_review(code: str) -> str: "round": 0, "max_rounds": 2, } - 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) + final = await app.ainvoke(state) + print("Final draft review:\n", final['draft_review']) + print("Scores:", final['criteria_scores']) + print("Verdict:", final['verdict']) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(demo())