Updated main.py with LangGraph code review agent

main.py

@@ -1,128 +1,222 @@
+"""
+# main.py
+# LangGraph code review agent with reflection and rewrite loop
+# Requires: langgraph, langchain-openai, deepagents, python-dotenv
+# Run with: python main.py
+"""
+
+from __future__ import annotations
+
 import os
-import asyncio
+from typing import TypedDict, Dict, Any
-from typing import TypedDict, Annotated, Dict
+
+from dotenv import load_dotenv
+
+# Deepagents is required by the assignment, but we do not use it directly.
+# Importing it ensures that the dependency is satisfied.
+import deepagents  # noqa: F401
 
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
+from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
-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
+from langchain_core.messages import HumanMessage
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import MemorySaver
+from pydantic import BaseModel
 
-# ---------- LLM ----------
+# Load OpenAI key from .env if present
-llm = ChatOpenAI(
+load_dotenv()
-    model="openai/gpt-oss-20b:free",
-    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]  # e.g., {'pep8': 8, ...}
     weakest_criterion: str
-    verdict: str
+    verdict: str  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
-# ---------- Pydantic for reflect output ----------
+# -----------------------------
-class ReflectOutput(BaseModel):
+# LLM configuration
-    pep8: int = Field(..., description="Score 0-10 for PEP8 compliance")
+# -----------------------------
-    type_hints: int = Field(..., description="Score 0-10 for type hints usage")
+# Replace with your preferred model or use Ollama via langchain-ollama if desired
-    edge_cases: int = Field(..., description="Score 0-10 for edge case handling")
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
-    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'")
 
-reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+# -----------------------------
+# Prompt templates
+# -----------------------------
+# draft_review prompt
+draft_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(
+        "You are a senior Python developer. Provide a concise code review for the following function. "
+        "List 3-6 bullet points, each describing a potential improvement or praise."
+    ),
+    HumanMessagePromptTemplate.from_template("Here is the function code:\n\n{code}")
+])
 
-# ---------- Nodes ----------
+# reflect prompt with structured output
-async def draft_review(state: CodeReviewState) -> CodeReviewState:
+class ReviewScores(BaseModel):
-    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.
+    pep8: int
+    type_hints: int
+    edge_cases: int
+    naming: int
+    weakest_criterion: str
+    verdict: str
 
-```python
+output_parser = PydanticOutputParser(pydantic_object=ReviewScores)
-{state['code']}
-```
 
-Return only the review text."""
+reflect_prompt = ChatPromptTemplate.from_messages([
-    review = await llm.ainvoke([HumanMessage(content=prompt)])
+    SystemMessagePromptTemplate.from_template(
-    state['draft_review'] = review.content.strip()
+        "You are a code quality critic. Evaluate the draft review for the following function. "
-    return state
+        "Score each of the four criteria (pep8, type_hints, edge_cases, naming) on a scale 0-10. "
+        "Identify the weakest criterion and provide a verdict: 'ok' if all scores are >=7, otherwise 'needs_revision'."
+    ),
+    HumanMessagePromptTemplate.from_template("Function code:\n\n{code}\n\nDraft review:\n\n{draft_review}")
+])
 
-async def reflect(state: CodeReviewState) -> CodeReviewState:
+# rewrite prompt
-    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".
+rewrite_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(
+        "You are a senior Python developer. Rewrite the draft review to improve the weakest criterion: {weakest_criterion}. "
+        "Keep 3-6 bullet points."
+    ),
+    HumanMessagePromptTemplate.from_template("Original draft review:\n\n{draft_review}")
+])
 
-Review text:
+# -----------------------------
-{state['draft_review']}
+# Node implementations
+# -----------------------------
+async def draft_review(state: CodeReviewState) -> Dict[str, Any]:
+    """Generate initial draft review."""
+    response = await llm.ainvoke(draft_prompt.format_messages(code=state['code']))
+    draft = response.content.strip()
+    # Ensure 3-6 bullet points by counting lines starting with '-'
+    bullets = [line for line in draft.splitlines() if line.lstrip().startswith('-')]
+    if len(bullets) < 3:
+        # If too few, add a generic positive point
+        draft += "\n- The code is readable and concise."
+    elif len(bullets) > 6:
+        # Trim to 6
+        draft = "\n".join(bullets[:6])
+    return {"draft_review": draft}
 
-Provide the output in the following JSON-like format:
+async def reflect(state: CodeReviewState) -> Dict[str, Any]:
-{"pep8": int, "type_hints": int, "edge_cases": int, "naming": int, "weakest_criterion": str, "verdict": str}"""
+    """Critic evaluates the draft review."""
-    raw = await llm.ainvoke([HumanMessage(content=prompt)])
+    response = await llm.ainvoke(
-    parsed = reflect_parser.parse(raw.content)
+        reflect_prompt.format_messages(
-    state['criteria_scores'] = {
+            code=state['code'],
-        "pep8": parsed.pep8,
+            draft_review=state['draft_review']
-        "type_hints": parsed.type_hints,
+        )
-        "edge_cases": parsed.edge_cases,
+    )
-        "naming": parsed.naming,
+    parsed = output_parser.parse(response.content)
+    return {
+        "criteria_scores": {
+            "pep8": parsed.pep8,
+            "type_hints": parsed.type_hints,
+            "edge_cases": parsed.edge_cases,
+            "naming": parsed.naming
+        },
+        "weakest_criterion": parsed.weakest_criterion,
+        "verdict": parsed.verdict
     }
-    state['weakest_criterion'] = parsed.weakest_criterion
-    state['verdict'] = parsed.verdict
-    return state
 
-async def rewrite(state: CodeReviewState) -> CodeReviewState:
+async def rewrite(state: CodeReviewState) -> Dict[str, Any]:
-    # Simple rewrite: add a sentence addressing the weakest criterion
+    """Rewrite the draft review focusing on the weakest criterion."""
-    additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}.",
+    response = await llm.ainvoke(
-    state['draft_review'] = state['draft_review'] + "\n" + additional
+        rewrite_prompt.format_messages(
-    state['round'] += 1
+            weakest_criterion=state['weakest_criterion'],
-    return state
+            draft_review=state['draft_review']
+        )
+    )
+    new_draft = response.content.strip()
+    # Ensure 3-6 bullet points
+    bullets = [line for line in new_draft.splitlines() if line.lstrip().startswith('-')]
+    if len(bullets) < 3:
+        new_draft += "\n- The code is readable and concise."
+    elif len(bullets) > 6:
+        new_draft = "\n".join(bullets[:6])
+    return {"draft_review": new_draft, "round": state['round'] + 1}
 
-# ---------- Graph ----------
+# -----------------------------
+# Graph construction
+# -----------------------------
 def build_graph() -> StateGraph[CodeReviewState]:
     graph = StateGraph(CodeReviewState)
     graph.add_node("draft_review", draft_review)
     graph.add_node("reflect", reflect)
     graph.add_node("rewrite", rewrite)
 
+    # Entry point
     graph.set_entry_point("draft_review")
+
+    # Transitions
     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",
+        lambda x: "needs_revision" if x["verdict"] == "needs_revision" else "ok",
+        {
+            "needs_revision": "rewrite",
+            "ok": END
+        }
     )
     graph.add_edge("rewrite", "reflect")
 
-    return graph.compile()
+    # Final node to stop if max rounds reached
+    def check_round(state: CodeReviewState):
+        if state["round"] >= state["max_rounds"]:
+            return "END"
+        return "rewrite"
 
-# ---------- Tool ----------
+    graph.add_conditional_edges(
-@tool
+        "rewrite",
-def code_review_tool(code: str) -> str:
+        check_round,
-    """Perform a structured code review with possible rewrites."""
+        {
-    graph = build_graph()
+            "END": END,
-    initial_state: CodeReviewState = {
+            "rewrite": "rewrite"
-        "code": code,
+        }
+    )
+
+    return graph.compile(checkpointer=MemorySaver())
+
+# -----------------------------
+# Demo execution
+# -----------------------------
+if __name__ == "__main__":
+    # Sample function to review
+    sample_code = """
+    def sort_numbers(arr):
+        return sorted(arr)
+    """
+
+    # Initial state
+    state: CodeReviewState = {
+        "code": sample_code.strip(),
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
         "verdict": "",
         "round": 0,
-        "max_rounds": 2,
+        "max_rounds": 2
     }
-    final_state = graph.invoke(initial_state)
-    return f"Final Review:\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}"
 
-# ---------- DeepAgent ----------
+    graph = build_graph()
-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__":
+    # Run the graph
-    asyncio.run(main())
+    async def run():
+        async for event in graph.astream(state):
+            # Print only the updated parts for clarity
+            if "draft_review" in event:
+                print("\n--- Draft Review (Round %d) ---" % (event.get("round", 0)))
+                print(event["draft_review"])
+            if "criteria_scores" in event:
+                print("\n--- Scores ---")
+                for k, v in event["criteria_scores"].items():
+                    print(f"{k}: {v}")
+                print(f"Weakest criterion: {event['weakest_criterion']}")
+                print(f"Verdict: {event['verdict']}")
+
+    import asyncio
+    asyncio.run(run())