Add main.py implementing LangGraph code review agent

main.py

@@ -1,105 +1,135 @@
-import asyncio
 import os
+import asyncio
+from typing import TypedDict, Dict
+
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
 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
 
 # LLM setup
 llm = ChatOpenAI(
-    model="openai/gpt-4o-mini",
+    model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
     api_key=os.getenv("OPENAI_API_KEY"),
     temperature=0.0,
 )
 
-# Pydantic models for reflection
+# State definition
-class CriteriaScores(BaseModel):
-    pep8: int
-    type_hints: int
-    edge_cases: int
-    naming: int
-
-class ReviewResult(BaseModel):
-    scores: CriteriaScores
-    weakest: str
-    verdict: str
-
-# Tool to run shell commands
-@tool
-def run_command(command: str) -> str:
-    """Execute a shell command and return its output."""
-    try:
-        result = os.popen(command).read()
-        return result.strip() or "(no output)"
-    except Exception as e:
-        return f"Error: {e}"
-
-# Draft review node
-async def draft_review(state: dict):
-    code = state["code"]
-    prompt = f"Write a concise code review for the following Python function. Provide 3-6 bullet points highlighting strengths and areas for improvement.\n\n{code}"
-    response = llm.invoke([HumanMessage(content=prompt)])
-    state["draft_review"] = response.content
-    return state
-
-# Reflect node
-async def reflect(state: dict):
-    review = state["draft_review"]
-    prompt = f"Score the following code review on 4 criteria: PEP8, type hints, edge cases, naming. Return JSON with keys pep8, type_hints, edge_cases, naming (0-10). Also provide the weakest criterion and verdict ('ok' if all >=7 else 'needs_revision').\n\n{review}"
-    response = llm.invoke([HumanMessage(content=prompt)])
-    try:
-        data = ReviewResult.parse_raw(response.content)
-    except Exception:
-        # fallback simple parse
-        data = ReviewResult(scores=CriteriaScores(pep8=5,type_hints=5,edge_cases=5,naming=5),weakest="pep8",verdict="needs_revision")
-    state["criteria_scores"] = data.scores.dict()
-    state["weakest_criterion"] = data.weakest
-    state["verdict"] = data.verdict
-    return state
-
-# Rewrite node
-async def rewrite(state: dict):
-    weakest = state["weakest_criterion"]
-    review = state["draft_review"]
-    prompt = f"Improve the code review focusing on the {weakest} aspect. Keep the rest unchanged.\n\n{review}"
-    response = llm.invoke([HumanMessage(content=prompt)])
-    state["draft_review"] = response.content
-    state["round"] += 1
-    return state
-
-# Graph definition
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
-    criteria_scores: dict
+    criteria_scores: Dict[str, int]
     weakest_criterion: str
     verdict: str
     round: int
     max_rounds: int
 
-workflow = StateGraph(CodeReviewState)
+# Reflect output model
-workflow.add_node("draft", draft_review)
+class ReflectOutput(BaseModel):
-workflow.add_node("reflect", reflect)
+    pep8: int = Field(description="Score 0-10 for PEP8 compliance")
-workflow.add_node("rewrite", rewrite)
+    type_hints: int = Field(description="Score 0-10 for type hints usage")
-workflow.add_conditional_edges(START, lambda _: "draft")
+    edge_cases: int = Field(description="Score 0-10 for edge case coverage")
-workflow.add_conditional_edges("draft", lambda _: "reflect")
+    naming: int = Field(description="Score 0-10 for naming conventions")
-workflow.add_conditional_edges("reflect", lambda s: "rewrite" if s["verdict"]=="needs_revision" and s["round"]<s["max_rounds"] else "END")
+    weakest_criterion: str = Field(description="Criterion with lowest score")
-workflow.add_conditional_edges("rewrite", lambda _: "reflect")
+    verdict: str = Field(description="'ok' or 'needs_revision'")
-graph = workflow.compile()
+
+reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+
+# Dummy tool for agent
+@tool
+def echo_tool(query: str) -> str:
+    return query
+
+# Agent creation
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+backend = CompositeBackend(
+    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
+    routes={},
+)
+agent = create_deep_agent(
+    model=llm,
+    tools=[echo_tool],
+    backend=backend,
+    system_prompt="You are a code review assistant.",
+)
+
+# Node functions
+async def draft_review(state: CodeReviewState) -> CodeReviewState:
+    prompt = f"""Write a concise code review (3-6 points) for the following Python function. Focus on style, correctness, and potential improvements.
+
+```python
+{state['code']}
+```
+
+Return only the review text."""
+    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "draft-review"}})
+    review_text = response["messages"][-1].content
+    state["draft_review"] = review_text
+    return state
+
+async def reflect(state: CodeReviewState) -> CodeReviewState:
+    prompt = f"""You are a senior reviewer. Evaluate the following review text against four criteria: PEP8, type hints, edge cases, naming. Assign each a score 0-10. Identify the weakest criterion and give a verdict: 'ok' if all scores >=7, else 'needs_revision'. Return a JSON with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.
+
+Review:
+{state['draft_review']}"""
+    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "reflect"}})
+    json_text = response["messages"][-1].content
+    try:
+        parsed = reflect_parser.parse(json_text)
+    except Exception:
+        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
+
+async def rewrite(state: CodeReviewState) -> CodeReviewState:
+    crit = state["weakest_criterion"]
+    prompt = f"""Improve the review section that addresses the weakest criterion '{crit}'. Provide a more detailed point for that criterion. Keep the rest of the review unchanged.
+
+Current review:
+{state['draft_review']}"""
+    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "rewrite"}})
+    new_review = response["messages"][-1].content
+    state["draft_review"] = new_review
+    state["round"] += 1
+    return state
+
+# Graph definition
+from langgraph.graph import StateGraph
+
+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(
+    "reflect",
+    lambda x: "END" if x["verdict"] == "ok" else "rewrite" if x["round"] < x["max_rounds"] else "END",
+)
+
+graph.add_edge("rewrite", "reflect")
+app = graph.compile()
 
 # Demo function
-async def run_demo():
+async def demo():
-    code = """
+    sample_code = """def sort_numbers(arr):
-def sort_numbers(arr):
+    return sorted(arr)"""
-    return sorted(arr)
+    init_state: CodeReviewState = {
-"""
+        "code": sample_code,
-    state: CodeReviewState = {
-        "code": code,
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
@@ -107,9 +137,18 @@ def sort_numbers(arr):
         "round": 0,
         "max_rounds": 2,
     }
-    final = await graph.ainvoke(state)
+    result = await app.ainvoke(init_state)
-    print("Final review:\n", final["draft_review"])
+    print("--- Draft Review ---")
-    print("Scores:", final["criteria_scores"])
+    print(result["draft_review"])
+    print("\n--- Scores ---")
+    print(result["criteria_scores"])
+    print("\n--- Verdict ---")
+    print(result["verdict"])
+    if result["verdict"] == "needs_revision":
+        print("\n--- Final Review After Rewrite ---")
+        print(result["draft_review"])
+        print("\n--- Final Scores ---")
+        print(result["criteria_scores"])
 
 if __name__ == "__main__":
-    asyncio.run(run_demo())
+    asyncio.run(demo())