diff --git a/main.py b/main.py index cd3f165..b65101c 100644 --- a/main.py +++ b/main.py @@ -1,105 +1,135 @@ -import asyncio import os +import asyncio +from typing import TypedDict, Dict + from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage from langchain.tools import tool 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 langchain_core.output_parsers import PydanticOutputParser # LLM setup llm = ChatOpenAI( - model="openai/gpt-4o-mini", + model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0, ) -# Pydantic models for reflection -class CriteriaScores(BaseModel): - pep8: int - type_hints: int - edge_cases: int - naming: int - -class ReviewResult(BaseModel): - scores: CriteriaScores - weakest: str - verdict: str - -# Tool to run shell commands -@tool -def run_command(command: str) -> str: - """Execute a shell command and return its output.""" - try: - result = os.popen(command).read() - return result.strip() or "(no output)" - except Exception as e: - return f"Error: {e}" - -# Draft review node -async def draft_review(state: dict): - code = state["code"] - prompt = f"Write a concise code review for the following Python function. Provide 3-6 bullet points highlighting strengths and areas for improvement.\n\n{code}" - response = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = response.content - return state - -# Reflect node -async def reflect(state: dict): - review = state["draft_review"] - prompt = f"Score the following code review on 4 criteria: PEP8, type hints, edge cases, naming. Return JSON with keys pep8, type_hints, edge_cases, naming (0-10). Also provide the weakest criterion and verdict ('ok' if all >=7 else 'needs_revision').\n\n{review}" - response = llm.invoke([HumanMessage(content=prompt)]) - try: - data = ReviewResult.parse_raw(response.content) - except Exception: - # fallback simple parse - data = ReviewResult(scores=CriteriaScores(pep8=5,type_hints=5,edge_cases=5,naming=5),weakest="pep8",verdict="needs_revision") - state["criteria_scores"] = data.scores.dict() - state["weakest_criterion"] = data.weakest - state["verdict"] = data.verdict - return state - -# Rewrite node -async def rewrite(state: dict): - weakest = state["weakest_criterion"] - review = state["draft_review"] - prompt = f"Improve the code review focusing on the {weakest} aspect. Keep the rest unchanged.\n\n{review}" - response = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = response.content - state["round"] += 1 - return state - -# Graph definition +# State definition class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: dict + criteria_scores: Dict[str, int] weakest_criterion: str verdict: str round: int max_rounds: int -workflow = StateGraph(CodeReviewState) -workflow.add_node("draft", draft_review) -workflow.add_node("reflect", reflect) -workflow.add_node("rewrite", rewrite) -workflow.add_conditional_edges(START, lambda _: "draft") -workflow.add_conditional_edges("draft", lambda _: "reflect") -workflow.add_conditional_edges("reflect", lambda s: "rewrite" if s["verdict"]=="needs_revision" and s["round"] str: + return query + +# Agent creation +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +backend = CompositeBackend( + default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), + routes={}, +) +agent = create_deep_agent( + model=llm, + tools=[echo_tool], + backend=backend, + system_prompt="You are a code review assistant.", +) + +# Node functions +async def draft_review(state: CodeReviewState) -> CodeReviewState: + prompt = f"""Write a concise code review (3-6 points) for the following Python function. Focus on style, correctness, and potential improvements. + +```python +{state['code']} +``` + +Return only the review text.""" + response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "draft-review"}}) + review_text = response["messages"][-1].content + state["draft_review"] = review_text + return state + +async def reflect(state: CodeReviewState) -> CodeReviewState: + prompt = f"""You are a senior reviewer. Evaluate the following review text against four criteria: PEP8, type hints, edge cases, naming. Assign each a score 0-10. Identify the weakest criterion and give a verdict: 'ok' if all scores >=7, else 'needs_revision'. Return a JSON with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict. + +Review: +{state['draft_review']}""" + response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "reflect"}}) + json_text = response["messages"][-1].content + try: + parsed = reflect_parser.parse(json_text) + except Exception: + 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 + +async def rewrite(state: CodeReviewState) -> CodeReviewState: + crit = state["weakest_criterion"] + prompt = f"""Improve the review section that addresses the weakest criterion '{crit}'. Provide a more detailed point for that criterion. Keep the rest of the review unchanged. + +Current review: +{state['draft_review']}""" + response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "rewrite"}}) + new_review = response["messages"][-1].content + state["draft_review"] = new_review + state["round"] += 1 + return state + +# Graph definition +from langgraph.graph import StateGraph + +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" else "rewrite" if x["round"] < x["max_rounds"] else "END", +) + +graph.add_edge("rewrite", "reflect") +app = graph.compile() # Demo function -async def run_demo(): - code = """ -def sort_numbers(arr): - return sorted(arr) -""" - state: CodeReviewState = { - "code": code, +async def demo(): + sample_code = """def sort_numbers(arr): + return sorted(arr)""" + init_state: CodeReviewState = { + "code": sample_code, "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -107,9 +137,18 @@ def sort_numbers(arr): "round": 0, "max_rounds": 2, } - final = await graph.ainvoke(state) - print("Final review:\n", final["draft_review"]) - print("Scores:", final["criteria_scores"]) + result = await app.ainvoke(init_state) + print("--- Draft Review ---") + print(result["draft_review"]) + print("\n--- Scores ---") + print(result["criteria_scores"]) + print("\n--- Verdict ---") + print(result["verdict"]) + if result["verdict"] == "needs_revision": + print("\n--- Final Review After Rewrite ---") + print(result["draft_review"]) + print("\n--- Final Scores ---") + print(result["criteria_scores"]) if __name__ == "__main__": - asyncio.run(run_demo()) + asyncio.run(demo())