diff --git a/main.py b/main.py index 40b96a9..84bc222 100644 --- a/main.py +++ b/main.py @@ -1,13 +1,15 @@ import os import asyncio -from typing import TypedDict, Annotated +from typing import TypedDict, Annotated, Dict from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, AIMessage -from langchain_core.output_parsers import PydanticOutputParser -from langchain_core.pydantic_v1 import BaseModel, Field +from langchain_core.messages import HumanMessage +from langchain.tools import tool + 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 # ---------- LLM ---------- llm = ChatOpenAI( @@ -21,124 +23,106 @@ llm = ChatOpenAI( class CodeReviewState(TypedDict): code: str draft_review: str - criteria_scores: dict[str, int] + criteria_scores: Dict[str, int] weakest_criterion: str - verdict: str # "ok" | "needs_revision" + verdict: str round: int max_rounds: int -# ---------- Structured output for reflect ---------- -class ReflectionOutput(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 handling edge cases") - naming: int = Field(description="Score 0-10 for naming conventions") - weakest_criterion: str = Field(description="Criterion with lowest score") - verdict: str = Field(description="'ok' or 'needs_revision'") +# ---------- Pydantic for reflect output ---------- +class ReflectOutput(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") + weakest_criterion: str = Field(..., description="Name of the weakest criterion") + verdict: str = Field(..., description="'ok' or 'needs_revision'") -parser = PydanticOutputParser(pydantic_object=ReflectionOutput) +reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput) # ---------- Nodes ---------- +async def draft_review(state: CodeReviewState) -> CodeReviewState: + prompt = f"""Please write a concise code review (3-6 bullet points) for the following Python function. Focus on style, type hints, edge cases, and naming. -def draft_review_node(state: CodeReviewState) -> CodeReviewState: - code = state["code"] - prompt = ( - "You are a senior Python developer.\n" - "Given the following function, write a concise code review (3-6 bullet points).\n" - "Focus on style, correctness, edge cases, and naming.\n" - f"Function:\n{code}\n\nReview:" # LLM will output review - ) - response = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = response.content +```python +{state['code']} +``` + +Return only the review text.""" + review = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft_review'] = review.content.strip() return state +async def reflect(state: CodeReviewState) -> CodeReviewState: + prompt = f"""You are a code review critic. Evaluate the following review text and assign scores 0-10 for each of the four criteria: pep8, type_hints, edge_cases, naming. Also identify the weakest criterion and decide if the review is "ok" or "needs_revision". -def reflect_node(state: CodeReviewState) -> CodeReviewState: - review = state["draft_review"] - code = state["code"] - prompt = ( - "You are an automated code review critic.\n" - "Given the code and its review, assign a score 0-10 for each of the following criteria:\n" - "- pep8: PEP8 compliance\n" - "- type_hints: use of type hints\n" - "- edge_cases: handling of edge cases\n" - "- naming: clarity of names\n" - "Return the scores, the weakest criterion, and a verdict ('ok' if all scores >=7, else 'needs_revision').\n" - f"Code:\n{code}\n\nReview:\n{review}\n\nOutput in JSON with fields: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict." - ) - response = llm.invoke([HumanMessage(content=prompt)]) - try: - out = parser.parse(response.content) - except Exception as e: - # Fallback: simple parsing if JSON is not strict - import json - out = json.loads(response.content) - state["criteria_scores"] = { - "pep8": out.pep8, - "type_hints": out.type_hints, - "edge_cases": out.edge_cases, - "naming": out.naming, +Review text: +{state['draft_review']} + +Provide the output in the following JSON-like format: +{"pep8": int, "type_hints": int, "edge_cases": int, "naming": int, "weakest_criterion": str, "verdict": str}""" + raw = await llm.ainvoke([HumanMessage(content=prompt)]) + parsed = reflect_parser.parse(raw.content) + state['criteria_scores'] = { + "pep8": parsed.pep8, + "type_hints": parsed.type_hints, + "edge_cases": parsed.edge_cases, + "naming": parsed.naming, } - state["weakest_criterion"] = out.weakest_criterion - state["verdict"] = out.verdict + state['weakest_criterion'] = parsed.weakest_criterion + state['verdict'] = parsed.verdict return state - -def rewrite_node(state: CodeReviewState) -> CodeReviewState: - weakest = state["weakest_criterion"] - review = state["draft_review"] - code = state["code"] - prompt = ( - "You are a senior Python developer tasked with improving a code review.\n" - f"The current review is:\n{review}\n\nThe weakest criterion is '{weakest}'.\n" - "Rewrite only the part of the review that addresses this criterion, making it stronger and more specific.\n" - "Keep the rest of the review unchanged.\n" - "Output only the updated review." - ) - response = llm.invoke([HumanMessage(content=prompt)]) - state["draft_review"] = response.content - state["round"] = state.get("round", 0) + 1 +async def rewrite(state: CodeReviewState) -> CodeReviewState: + # Simple rewrite: add a sentence addressing the weakest criterion + additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}." + state['draft_review'] = state['draft_review'] + "\n" + additional + state['round'] += 1 return state # ---------- Graph ---------- -builder = StateGraph(CodeReviewState) -builder.add_node("draft_review", draft_review_node) -builder.add_node("reflect", reflect_node) -builder.add_node("rewrite", rewrite_node) +def build_graph() -> StateGraph[CodeReviewState]: + graph = StateGraph(CodeGraphState) + 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( - "reflect", - lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END", -) -builder.add_edge("rewrite", "reflect") -builder.add_edge("END", END) + graph.set_entry_point("draft_review") + graph.add_edge("draft_review", "reflect") + graph.add_conditional_edges( + "reflect", + lambda x: "END" if x['verdict'] == "ok" or x['round'] >= x['max_rounds'] else "rewrite", + ) + graph.add_edge("rewrite", "reflect") -graph = builder.compile() + return graph.compile() -# ---------- Demo ---------- -async def main(): - demo_code = """ - def sort_numbers(arr): - return sorted(arr) - """ +# ---------- Tool ---------- +@tool +def code_review_tool(code: str) -> str: + """Perform a structured code review with possible rewrites.""" + graph = build_graph() initial_state: CodeReviewState = { - "code": demo_code.strip(), - "draft_review": "", # will be filled + "code": code, + "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", "verdict": "", "round": 0, "max_rounds": 2, } - result = await graph.ainvoke(initial_state) - print("\n--- Final Review ---") - print(result["draft_review"]) - print("\n--- Scores ---") - for k, v in result["criteria_scores"].items(): - print(f"{k}: {v}") - print(f"Verdict: {result['verdict']}") + final_state = graph.invoke(initial_state) + return f"Final Review:\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}" + +# ---------- DeepAgent ---------- +async def main(): + sample_code = """ + def sort_numbers(arr): + return sorted(arr) + """ + # Use the code_review_tool directly without deepagents + final_review = code_review_tool(sample_code) + print(final_review) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(main()) \ No newline at end of file