diff --git a/main.py b/main.py new file mode 100644 index 0000000..cd3f165 --- /dev/null +++ b/main.py @@ -0,0 +1,115 @@ +import asyncio +import os +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", + 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 +class CodeReviewState(TypedDict): + code: str + draft_review: str + criteria_scores: dict + 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"]