diff --git a/main.py b/main.py index 191d774..55fdfaf 100644 --- a/main.py +++ b/main.py @@ -1,96 +1,146 @@ import os -import asyncio +import json +from typing import TypedDict, Dict +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_agent from langchain_openai import ChatOpenAI -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 langchain_core.output_parsers import PydanticOutputParser +from langchain_core.messages import HumanMessage, AIMessage -# LLM setup -llm = ChatOpenAI( - model="gpt-4o-mini", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=0.0, -) - -# State definition +# --------------------- +# 1. State definition +# --------------------- class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: dict[str, int] + criteria_scores: Dict[str, int] # {"pep8": 0-10, "type_hints": 0-10, "edge_cases": 0-10, "naming": 0-10} weakest_criterion: str - verdict: str + verdict: str # "ok" | "needs_revision" round: int max_rounds: int -# 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 +# --------------------- +# 2. LLM setup +# --------------------- +# Use OpenAI or Ollama based on env variable +if os.getenv("USE_OLLAMA", "false").lower() == "true": + from langchain_ollama import ChatOllama + llm = ChatOllama(model="llama3", temperature=0.2) +else: + llm = ChatOpenAI(temperature=0.2, model_name="gpt-4o-mini") -# Node: reflect -class ReflectOutput(BaseModel): - scores: dict[str, int] - weakest: str - verdict: str +# --------------------- +# 3. Node definitions +# --------------------- -parser = PydanticOutputParser(pydantic_object=ReflectOutput) +def draft_review_fn(state: CodeReviewState) -> Dict: + code = state["code"] + prompt = f""" +You are a senior Python developer. Provide a concise code review (3-6 bullet points) for the following function. Focus on style, correctness, and potential improvements. -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 +Function: +{code} -# 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 +Review: +""" + response = llm.invoke([HumanMessage(content=prompt)]) + review = response.content.strip() + return {"draft_review": review} -# Graph -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( +def reflect_fn(state: CodeReviewState) -> Dict: + review = state["draft_review"] + prompt = f""" +You are an automated code review critic. Evaluate the following code review on four criteria: PEP8 compliance, type hints usage, edge case handling, and naming conventions. Assign each a score from 0 to 10. Also determine the weakest criterion and a verdict: "ok" if all scores are 7 or higher, otherwise "needs_revision". + +Review: +{review} + +Respond in JSON with keys: "pep8", "type_hints", "edge_cases", "naming", "weakest_criterion", "verdict". +""" + response = llm.invoke([HumanMessage(content=prompt)]) + try: + data = json.loads(response.content) + except Exception: + # Fallback: simple parsing + data = { + "pep8": 5, + "type_hints": 5, + "edge_cases": 5, + "naming": 5, + "weakest_criterion": "pep8", + "verdict": "needs_revision" + } + return { + "criteria_scores": { + "pep8": int(data.get("pep8", 0)), + "type_hints": int(data.get("type_hints", 0)), + "edge_cases": int(data.get("edge_cases", 0)), + "naming": int(data.get("naming", 0)) + }, + "weakest_criterion": data.get("weakest_criterion", "pep8"), + "verdict": data.get("verdict", "needs_revision") + } + + +def rewrite_fn(state: CodeReviewState) -> Dict: + weakest = state["weakest_criterion"] + review = state["draft_review"] + prompt = f""" +You are a senior Python developer. The following code review has been identified as weak in the "{weakest}" criterion. Rewrite only the part of the review that addresses this criterion, improving it significantly. Keep the rest of the review unchanged. + +Original Review: +{review} + +Rewritten Review: +""" + response = llm.invoke([HumanMessage(content=prompt)]) + new_review = response.content.strip() + return {"draft_review": new_review, "round": state["round"] + 1} + +# --------------------- +# 4. Graph construction +# --------------------- +builder = StateGraph(CodeReviewState) + +builder.add_node("draft_review", draft_review_fn) +builder.add_node("reflect", reflect_fn) +builder.add_node("rewrite", rewrite_fn) + +# Entry point +builder.set_entry_point("draft_review") + +# Transitions +builder.add_edge("draft_review", "reflect") +builder.add_conditional_edges( "reflect", - lambda s: "rewrite" if s['verdict']=='needs_revision' and s['round']