diff --git a/main.py b/main.py index 478877d..9485a5a 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,37 @@ 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 -# ---------- State definition ---------- +# Load environment variables +from dotenv import load_dotenv +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 +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) + +# State definition class CodeReviewState(TypedDict): code: str draft_review: str @@ -24,129 +41,100 @@ class CodeReviewState(TypedDict): round: int max_rounds: int -# ---------- Structured output for critic ---------- -class CriticOutput(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)") - verdict: str = Field(description='Verdict: "ok" or "needs_revision"') +# Structured output for reflect node +class ReviewScores(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") + verdict: str = Field(description='\"ok\" or \"needs_revision\"') -critic_parser = PydanticOutputParser(pydantic_object=CriticOutput) +review_parser = PydanticOutputParser(pydantic_object=ReviewScores) -# ---------- 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"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 +# Draft review node +def draft_review(state: CodeReviewState) -> CodeReviewState: + prompt = f"""Write a concise code review for the following Python function. Provide 3-6 bullet points. Do not include any code blocks. The function is: +```python +{state["code"]} +```""" + response = llm.invoke(prompt) + state["draft_review"] = response.content.strip() print("\n--- Draft Review ---") - print(review) + print(state["draft_review"]) return state -async def reflect(state: CodeReviewState) -> CodeReviewState: - prompt = ( - 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}" - ) - 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 LLM output is not perfectly formatted - parsed = CriticOutput( - pep8=0, - type_hints=0, - edge_cases=0, - naming=0, - verdict="needs_revision", - ) - scores = { +# Reflect node +def reflect(state: CodeReviewState) -> CodeReviewState: + prompt = f"""You are a code review critic. Evaluate the following draft review for the given code. Provide scores 0-10 for each criterion and a verdict. Return JSON matching the schema: +{review_parser.get_format_instructions()} +Draft review: +{state["draft_review"]} + +Code: +```python +{state["code"]} +```""" + response = llm.invoke(prompt) + parsed = review_parser.parse(response.content) + state["criteria_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 + state["weakest_criterion"] = min(state["criteria_scores"], key=state["criteria_scores"].get) print("\n--- Reflection ---") - print(f"Scores: {scores}") - print(f"Weakest criterion: {weakest}") - print(f"Verdict: {parsed.verdict}") + print(f"Scores: {state['criteria_scores']}") + print(f"Weakest criterion: {state['weakest_criterion']}") + print(f"Verdict: {state['verdict']}") return state -async def rewrite(state: CodeReviewState) -> CodeReviewState: - prompt = ( - 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']}" - ) - result = await agent.ainvoke([HumanMessage(content=prompt)]) - new_review = result["messages"][-1].content.strip() - state["draft_review"] = new_review +# Rewrite node +def rewrite(state: CodeReviewState) -> CodeReviewState: + prompt = f"""Rewrite the section of the draft review that addresses the weakest criterion ({state['weakest_criterion']}) to improve it. Keep other parts unchanged. Return only the updated draft review. +Original draft review: +{state['draft_review']} + +Code: +```python +{state['code']} +```""" + response = llm.invoke(prompt) + state["draft_review"] = response.content.strip() state["round"] += 1 - print("\n--- Rewritten Review ---") - print(new_review) + print("\n--- Rewrite ---") + print(state["draft_review"]) return state -# ---------- Graph ---------- +# Graph definition def build_graph() -> StateGraph: graph = StateGraph(CodeReviewState) graph.add_node("draft_review", draft_review) graph.add_node("reflect", reflect) graph.add_node("rewrite", rewrite) - graph.add_edge(START, "draft_review") + graph.set_entry_point("draft_review") graph.add_edge("draft_review", "reflect") - def reflect_cond(state: CodeReviewState): + def condition(state: CodeReviewState): if state["verdict"] == "ok": - return END - if state["round"] < state["max_rounds"]: - return "rewrite" - return END + return "END" + if state["round"] >= state["max_rounds"]: + return "END" + return "rewrite" - graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", END: END}) + graph.add_conditional_edges("reflect", condition, ["rewrite", "END"]) graph.add_edge("rewrite", "reflect") return graph -# ---------- Demo ---------- -async def main(): - # Sample function to review - code_str = """def sort_numbers(arr): - return sorted(arr)""" - +# Tool that runs the graph +@tool +def run_review(code: str) -> str: + """Run a code review cycle on the provided Python function.""" initial_state: CodeReviewState = { - "code": code_str, + "code": code, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -154,16 +142,32 @@ async def main(): "round": 0, "max_rounds": 2, } - graph = build_graph() - app = graph.compile() - final_state = await app.ainvoke(initial_state) + final_state = graph.invoke(initial_state) + output = f"Final review after {final_state['round']} round(s):\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}\nVerdict: {final_state['verdict']}" + return output - print("\n=== Final State ===") - print(f"Round: {final_state['round']}") - print(f"Verdict: {final_state['verdict']}") - print(f"Draft Review:\n{final_state['draft_review']}") - print(f"Scores: {final_state['criteria_scores']}") +# Deepagents agent +agent = create_deep_agent( + model=llm, + tools=[run_review], + backend=backend, + system_prompt="You are a helpful code review agent.", +) + +# Demo function +demo_code = """ +def sort_numbers(arr): + return sorted(arr) +""" + +async def main(): + result = await agent.ainvoke( + {"messages": [HumanMessage(content=demo_code)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print("\n=== Final Output ===") + print(result["messages"][-1].content) if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file