task-6a22c713fd30e81cf315ea04/main.py

import os
import asyncio
from typing import TypedDict, Annotated, Dict

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

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ---------- Backend ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ---------- State ----------
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: Dict[str, int]
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int

# ---------- Pydantic models for structured output ----------
class ReviewScores(BaseModel):
    pep8: int = Field(..., ge=0, le=10)
    type_hints: int = Field(..., ge=0, le=10)
    edge_cases: int = Field(..., ge=0, le=10)
    naming: int = Field(..., ge=0, le=10)
    weakest_criterion: str
    verdict: str

class ReviewRewrite(BaseModel):
    draft_review: str
    criteria_scores: Dict[str, int]
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int

# ---------- Output parsers ----------
review_parser = PydanticOutputParser(pydantic_object=ReviewScores)
rewrite_parser = PydanticOutputParser(pydantic_object=ReviewRewrite)

# ---------- Nodes ----------

def draft_review_node(state: CodeReviewState) -> CodeReviewState:
    code = state["code"]
    prompt = f"""
You are a senior Python reviewer. Provide a concise code review for the following function. Output exactly 3-6 bullet points, each starting with a dash. Do not include any additional text.

{code}
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = response.content.strip()
    return state

# DESIGN DECISION: reflect node returns structured JSON with scores and verdict
# NECESSITY: required by assignment to have structured output for automated parsing
# OPTIMALITY: eliminates ambiguity and parsing errors compared to free text
# ALTERNATIVES CONSIDERED: free text parsing, regex extraction – rejected due to unreliability

def reflect_node(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""
You are an automated code quality critic. Evaluate the following draft review against these criteria:
- PEP8 compliance
- Presence of type hints
- Handling of edge cases
- Naming conventions

Return a JSON object with integer scores 0-10 for each criterion, the name of the weakest criterion, and a verdict "ok" or "needs_revision".

Draft review:
{state["draft_review"]}
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    parsed = review_parser.parse(response.content)
    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

# DESIGN DECISION: rewrite node focuses only on weakest criterion
# NECESSITY: assignment specifies targeted rewrite
# OPTIMALITY: keeps changes minimal and focused, avoids over‑engineering
# ALTERNATIVES CONSIDERED: full rewrite of review – rejected for unnecessary complexity

def rewrite_node(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""
You are a code reviewer. The previous draft review was:
{state["draft_review"]}

The weakest criterion is {state["weakest_criterion"]}. Rewrite only the part of the review that addresses this criterion, improving it. Keep the rest of the review unchanged. Output the updated draft review and updated scores (same format as in reflect). Also increment the round counter.
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    parsed = rewrite_parser.parse(response.content)
    state["draft_review"] = parsed.draft_review
    state["criteria_scores"] = parsed.criteria_scores
    state["weakest_criterion"] = parsed.weakest_criterion
    state["verdict"] = parsed.verdict
    state["round"] = parsed.round
    state["max_rounds"] = parsed.max_rounds
    return state

# ---------- Graph ----------

graph = StateGraph(CodeReviewState)

graph.add_node("draft_review", draft_review_node)
graph.add_node("reflect", reflect_node)
graph.add_node("rewrite", rewrite_node)

# Entry point
graph.add_edge(START, "draft_review")
graph.add_edge("draft_review", "reflect")
# Conditional edges after reflect

def decide_next(state: CodeReviewState):
    if state["verdict"] == "ok":
        return END
    if state["round"] < state["max_rounds"]:
        return "rewrite"
    return END

graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", END: END})
# After rewrite go back to reflect
graph.add_edge("rewrite", "reflect")

app = graph.compile()

# ---------- DeepAgent wrapper ----------
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a code review assistant.",
)

# ---------- CLI Demo ----------
async def main():
    # Example function to review
    code = """
    def sort_numbers(arr):
        return sorted(arr)
    """
    initial_state: CodeReviewState = {
        "code": code.strip(),
        "draft_review": "",  # will be filled
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = await app.ainvoke(initial_state)
    print("\n--- Final Review ---")
    print(result["draft_review"])
    print("\nScores:", result["criteria_scores"])
    print("Verdict:", result["verdict"])

if __name__ == "__main__":
    asyncio.run(main())