Add main.py

main.py

@@ -1,136 +1,84 @@
 import os
 import asyncio
-from typing import TypedDict, Annotated, Dict
-from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, SystemMessage
-from langchain_core.output_parsers import PydanticOutputParser
+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 deepagents import create_deep_agent
+from langchain_core.output_parsers import PydanticOutputParser
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from deepagents.tools import tool
 
-# ---------- LLM ----------
+# LLM setup
 llm = ChatOpenAI(
-    model="openai/gpt-oss-20b:free",
+    model="gpt-4o-mini",
     base_url="https://openrouter.ai/api/v1",
     api_key=os.getenv("OPENAI_API_KEY"),
     temperature=0.0,
 )
 
-# ---------- State ----------
+# State definition
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
-    criteria_scores: Dict[str, int]
+    criteria_scores: dict[str, int]
     weakest_criterion: str
     verdict: str
     round: int
     max_rounds: int
 
-# ---------- Structured output for reflect ----------
+# 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
+
+# Node: reflect
 class ReflectOutput(BaseModel):
-    pep8: int = Field(..., ge=0, le=10)
+    scores: dict[str, int]
-    type_hints: int = Field(..., ge=0, le=10)
+    weakest: str
-    edge_cases: int = Field(..., ge=0, le=10)
+    verdict: str
-    naming: int = Field(..., ge=0, le=10)
-    weakest_criterion: str = Field(...)
-    verdict: str = Field(..., regex="^(ok|needs_revision)$")
 
-reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+parser = PydanticOutputParser(pydantic_object=ReflectOutput)
 
-# ---------- Nodes ----------
+async def reflect(state: CodeReviewState):
-@tool
+    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']}"
-def draft_review(state: CodeReviewState) -> CodeReviewState:
+    msg = await llm.ainvoke([HumanMessage(content=prompt)])
-    """Generate an initial code review."""
+    out = parser.parse(msg.content)
-    prompt = (
+    state['criteria_scores'] = out.scores
-        "You are a senior Python reviewer.\n"
+    state['weakest_criterion'] = out.weakest
-        "Given the following function, write a concise code review (3–6 points).\n"
+    state['verdict'] = out.verdict
-        "Focus on style, correctness, and potential improvements.\n"
-        "Return only the review text.\n\n"
-        f"Function:\n{state['code']}"
-    )
-    review = llm.invoke([HumanMessage(content=prompt)]).content
-    state['draft_review'] = review
     return state
 
-@tool
+# Node: rewrite
-def reflect(state: CodeReviewState) -> CodeReviewState:
+async def rewrite(state: CodeReviewState):
-    """Critique the draft review and score four criteria."""
+    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']}"
-    prompt = (
+    msg = await llm.ainvoke([HumanMessage(content=prompt)])
-        "You are an automated code review critic.\n"
+    state['draft_review'] = msg.content
-        "Given the original code and the draft review, assign a score 0–10 for each of the following criteria:\n"
-        "- pep8: adherence to PEP8 style guide\n"
-        "- type_hints: use of type hints\n"
-        "- edge_cases: handling of edge cases\n"
-        "- naming: clarity of identifiers\n"
-        "Also identify the weakest criterion and decide if the review is "ok" or "needs_revision".\n"
-        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.\n"
-        "Do not include any other text.\n\n"
-        f"Code:\n{state['code']}\n\n"
-        f"Draft Review:\n{state['draft_review']}"
-    )
-    raw = llm.invoke([HumanMessage(content=prompt)]).content
-    try:
-        parsed = reflect_parser.parse(raw)
-    except Exception as e:
-        # Fallback: simple extraction
-        parsed = ReflectOutput(pep8=5, type_hints=5, edge_cases=5, naming=5, weakest_criterion="pep8", 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
-    return state
-
-@tool
-def rewrite(state: CodeReviewState) -> CodeReviewState:
-    """Rewrite the section of the draft review that addresses the weakest criterion."""
-    prompt = (
-        "You are a senior Python reviewer.\n"
-        "The draft review below has been critiqued. The weakest criterion is {criterion}.\n"
-        "Rewrite only the part of the review that addresses this criterion, improving it.\n"
-        "Keep the rest of the review unchanged.\n"
-        "Return the full updated review.\n\n"
-        f"Weakest criterion: {state['weakest_criterion']}\n\n"
-        f"Draft Review:\n{state['draft_review']}"
-    ).format(criterion=state['weakest_criterion'])
-    updated = llm.invoke([HumanMessage(content=prompt)]).content
-    state['draft_review'] = updated
     state['round'] += 1
     return state
 
-# ---------- Graph ----------
+# Graph
-builder = StateGraph(CodeReviewState)
+graph = StateGraph(CodeReviewState)
-builder.add_node("draft_review", draft_review)
+graph.add_node("draft_review", draft_review)
-builder.add_node("reflect", reflect)
+graph.add_node("reflect", reflect)
-builder.add_node("rewrite", rewrite)
+graph.add_node("rewrite", rewrite)
 
-builder.set_entry_point("draft_review")
+graph.set_entry_point("draft_review")
-builder.add_edge("draft_review", "reflect")
+graph.add_edge("draft_review", "reflect")
-builder.add_conditional_edges(
+graph.add_conditional_edges(
     "reflect",
-    lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END",
+    lambda s: "rewrite" if s['verdict']=='needs_revision' and s['round']<s['max_rounds'] else END,
 )
-builder.add_edge("rewrite", "reflect")
+graph.add_edge("rewrite", "reflect")
-builder.add_edge("END", END)
 
-graph = builder.compile()
+app = graph.compile()
 
-# ---------- DeepAgent wrapper ----------
+# Demo function
-backend = CompositeBackend([
+async def demo():
-    LocalShellBackend(workspace_dir="./workspace"),
+    code = """def sort_numbers(arr):
-    FilesystemBackend(),
+    return sorted(arr)"""
-])
+    state: CodeReviewState = {
-
-@tool
-def run_review(code: str) -> str:
-    """Run the LangGraph code review pipeline on the provided code."""
-    initial_state: CodeReviewState = {
         "code": code,
         "draft_review": "",
         "criteria_scores": {},
@@ -139,31 +87,10 @@ def run_review(code: str) -> str:
         "round": 0,
         "max_rounds": 2,
     }
-    final_state = graph.invoke(initial_state)
+    final = await app.ainvoke(state)
-    return (
+    print("Final draft review:\n", final['draft_review'])
-        f"Initial Draft Review:\n{final_state['draft_review']}\n\n"
+    print("Scores:", final['criteria_scores'])
-        f"Scores: {final_state['criteria_scores']}\n"
+    print("Verdict:", final['verdict'])
-        f"Verdict: {final_state['verdict']}\n"
-        f"Rounds: {final_state['round']}\n"
-    )
-
-agent = create_deep_agent(
-    model=llm,
-    tools=[run_review],
-    backend=backend,
-    system_prompt="You are a code review assistant.",
-)
-
-async def main():
-    code_example = """
-    def sort_numbers(arr):
-        return sorted(arr)
-    """
-    result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=f"Please review this code:\n{code_example}")]},
-        {"configurable": {"thread_id": "session-1"}},
-    )
-    print(result["messages"][-1].content)
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(demo())