diff --git a/main.py b/main.py index eabdb26..107f02e 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,21 @@ import os import asyncio -from typing import TypedDict, Annotated, Dict +from typing import TypedDict, Dict +from dotenv import load_dotenv 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 +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend -# ---------- LLM ---------- +# Load API key from .env +load_dotenv() + +# LLM configuration - OpenRouter llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -22,124 +23,140 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Backend ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +# Backend for deepagents - simple filesystem +backend = FilesystemBackend() -# ---------- State ---------- +# Create a deepagents agent that will be used inside the graph nodes +agent = create_deep_agent( + model=llm, + tools=[], + backend=backend, + system_prompt="You are a code review assistant.", +) + +# ---------- State definition ---------- class CodeReviewState(TypedDict): code: str draft_review: str criteria_scores: Dict[str, int] weakest_criterion: str - verdict: str + verdict: str # "ok" | "needs_revision" round: int max_rounds: int -# ---------- Pydantic model for reflect output ---------- -class ReflectOutput(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)") - weakest_criterion: str = Field(description="The criterion with the lowest score") - verdict: str = Field(description="'ok' or 'needs_revision'") +# ---------- Structured output for critic ---------- +class CriticOutput(BaseModel): + scores: Dict[str, int] = Field( + description="Scores for each criterion: pep8, type_hints, edge_cases, naming. Values 0-10." + ) + verdict: str = Field( + description='Verdict: "ok" if all scores >= 7, otherwise "needs_revision".' + ) -parser = PydanticOutputParser(pydantic_object=ReflectOutput) +critic_parser = PydanticOutputParser(pydantic_object=CriticOutput) -# ---------- Nodes ---------- +# ---------- Graph nodes ---------- async def draft_review(state: CodeReviewState) -> CodeReviewState: - prompt = f""" - Write a concise code review for the following Python function. Provide 3-6 bullet points. - - Function: - ```python - {state['code']} - ``` - """ - response = await llm.ainvoke(HumanMessage(content=prompt)) - state['draft_review'] = response.content.strip() + prompt = ( + f"Write a concise code review (3-6 bullet points) for the following Python function:\n\n" + f"{state['code']}\n\n" + "Focus on style, correctness, and potential improvements." + ) + response = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)]}, + {"configurable": {"thread_id": "draft_review"}}, + ) + review_text = response["messages"][-1].content.strip() + state["draft_review"] = review_text + print("\n--- Draft Review ---") + print(review_text) return state async def reflect(state: CodeReviewState) -> CodeReviewState: - prompt = f""" - Evaluate the following code and draft review. Score each of the four criteria on a scale 0-10: - - PEP8 compliance - - Type hints - - Edge case handling - - Naming conventions - - Provide the scores, identify the weakest criterion, and give a verdict ('ok' or 'needs_revision'). - - Code: - ```python - {state['code']} - ``` - - Draft Review: - ```text - {state['draft_review']} - ``` - - Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict. - """ - response = await llm.ainvoke(HumanMessage(content=prompt)) + prompt = ( + f"Evaluate the following code review and assign scores (0-10) for each criterion:\n\n" + f"Review:\n{state['draft_review']}\n\n" + "Criteria:\n" + "1. pep8: adherence to PEP8 style guide.\n" + "2. type_hints: presence and correctness of type hints.\n" + "3. edge_cases: handling of edge cases and robustness.\n" + "4. naming: clarity and consistency of names.\n\n" + "Return a JSON object with keys 'scores' (dict) and 'verdict' ('ok' or 'needs_revision')." + ) + response = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)]}, + {"configurable": {"thread_id": "reflect"}}, + ) + raw_output = response["messages"][-1].content.strip() try: - parsed = parser.parse(response.content) + parsed = critic_parser.parse(raw_output) except Exception as e: - # Fallback: set all scores to 0 and verdict to needs_revision - parsed = ReflectOutput(pep8=0, type_hints=0, edge_cases=0, naming=0, weakest_criterion="unknown", 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 + # Fallback: simple parsing if JSON is malformed + import json + parsed = CriticOutput(**json.loads(raw_output)) + state["criteria_scores"] = parsed.scores + # Determine weakest criterion + weakest = min(parsed.scores.items(), key=lambda kv: kv[1])[0] + state["weakest_criterion"] = weakest + state["verdict"] = parsed.verdict + print("\n--- Critic Scores ---") + for crit, score in parsed.scores.items(): + print(f"{crit}: {score}") + print(f"Weakest criterion: {weakest}") + print(f"Verdict: {parsed.verdict}") return state async def rewrite(state: CodeReviewState) -> CodeReviewState: - prompt = f""" - Rewrite the section of the draft review that addresses the weakest criterion: {state['weakest_criterion']}. - Keep all other parts of the review unchanged. - - Current Draft Review: - ```text - {state['draft_review']} - ``` - """ - response = await llm.ainvoke(HumanMessage(content=prompt)) - state['draft_review'] = response.content.strip() - state['round'] += 1 + state["round"] += 1 + prompt = ( + f"Rewrite the part of the review that addresses the weakest criterion " + f"('{state['weakest_criterion']}') to improve it. Keep the rest of the review unchanged.\n\n" + f"Original Review:\n{state['draft_review']}\n\n" + "Provide only the updated review." + ) + response = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)]}, + {"configurable": {"thread_id": "rewrite"}}, + ) + new_review = response["messages"][-1].content.strip() + state["draft_review"] = new_review + print("\n--- Rewritten Review (Round {}) ---".format(state["round"])) + print(new_review) return state -# ---------- Graph ---------- -review_graph = StateGraph(CodeReviewState) -review_graph.add_node("draft_review", draft_review) -review_graph.add_node("reflect", reflect) -review_graph.add_node("rewrite", rewrite) +# ---------- Graph construction ---------- +builder = StateGraph(CodeReviewState) -review_graph.set_entry_point("draft_review") -review_graph.add_edge("draft_review", "reflect") +builder.add_node("draft_review", draft_review) +builder.add_node("reflect", reflect) +builder.add_node("rewrite", rewrite) + +builder.add_edge(START, "draft_review") +builder.add_edge("draft_review", "reflect") # Conditional edges after reflect -review_graph.add_conditional_edges( - "reflect", - lambda state: "END" if state["verdict"] == "ok" else "rewrite" if state["round"] < state["max_rounds"] else "END", -) -review_graph.add_edge("rewrite", "reflect") +def reflect_conditional(state: CodeReviewState): + if state["verdict"] == "ok": + return END + if state["round"] < state["max_rounds"]: + return "rewrite" + return END -compiled_graph = review_graph.compile() +builder.add_conditional_edges("reflect", reflect_conditional) -# ---------- Tool ---------- -@tool -async def run_code_review(code: str) -> Dict: - """Run a code review on the provided Python function.""" +builder.add_edge("rewrite", "reflect") + +graph = builder.compile() + +# ---------- Demo ---------- +async def main(): + # Sample function to review + sample_code = """ +def sort_numbers(arr): + return sorted(arr) +""" initial_state: CodeReviewState = { - "code": code, + "code": sample_code.strip(), "draft_review": "", "criteria_scores": {}, "weakest_criterion": "", @@ -147,37 +164,13 @@ async def run_code_review(code: str) -> Dict: "round": 0, "max_rounds": 2, } - final_state = compiled_graph.invoke(initial_state) - return { - "final_review": final_state["draft_review"], - "scores": final_state["criteria_scores"], - "weakest_criterion": final_state["weakest_criterion"], - "verdict": final_state["verdict"], - "rounds": final_state["round"], - } -# ---------- Agent ---------- -agent = create_deep_agent( - model=llm, - tools=[run_code_review], - backend=backend, - system_prompt="You are a helpful code review agent.", -) - -# ---------- Demo ---------- -async def main(): - # Sample function to review - code = """ - def sort_numbers(arr): - return sorted(arr) - """ - result = await run_code_review(code) - print("\n=== Final Review ===\n") - print(result["final_review"]) - print("\n=== Scores ===\n") - for k, v in result["scores"].items(): - print(f"{k}: {v}") - print(f"\nVerdict: {result['verdict']} (Rounds: {result['rounds']})") + final_state = await graph.ainvoke(initial_state) + print("\n=== Final State ===") + print(f"Verdict: {final_state['verdict']}") + print(f"Rounds performed: {final_state['round']}") + print("\nFinal Review:") + print(final_state["draft_review"]) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(main()) \ No newline at end of file