Solution published: update main.py

main.py

@@ -1,31 +1,12 @@
-"""LangGraph Code Review Agent
-
-This repository implements a LangGraph agent that takes a Python function as input and produces a code review. The review is evaluated by a critic node that scores it on four criteria:
-
-1. PEP8 compliance
-2. Type hints
-3. Edge case handling
-4. Naming conventions
-
-If the critic returns "needs_revision" the rewrite node improves the weakest part of the review. The process repeats up to ``max_rounds`` times.
-
-The implementation uses only the technologies specified in the assignment: ``langgraph`` and ``langchain-ollama`` (or ``langchain-openai`` if you prefer). No vector database is used.
-
-Run the demo with ``python main.py``.
-"""
-
-from __future__ import annotations
-
 from typing import TypedDict, Dict
+import inspect
 
 from langgraph.graph import StateGraph, END
-from langgraph.prebuilt import create_structured_output_node
+from langchain_ollama import Ollama
-from langchain_ollama import ChatOllama
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StructuredOutputParser
+from langchain_core.output_parsers import JsonOutputParser
-from langchain_core.messages import HumanMessage
 
-# ---------- State ----------
+# Define the state
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
@@ -35,107 +16,81 @@ class CodeReviewState(TypedDict):
     round: int
     max_rounds: int
 
-# ---------- LLM ----------
+# LLM instance (Ollama)
-# Use Ollama; adjust model name if needed
+llm = Ollama(model="llama3.1")
-llm = ChatOllama(model="llama3")
 
-# ---------- Draft Review Node ----------
+# Node: draft_review
-DRAFT_PROMPT = ChatPromptTemplate.from_messages([
-    ("system", "You are a senior Python developer. Your task is to write a concise code review for the following function. Provide 3-6 points, each starting with a dash.")
-])
 
-async def draft_review(state: CodeReviewState) -> Dict[str, str]:
+def draft_review(state: CodeReviewState) -> Dict[str, str]:
-    prompt = DRAFT_PROMPT.format_messages(code=state["code"])
+    prompt = ChatPromptTemplate.from_messages([
-    response = await llm.ainvoke(prompt)
+        ("system", "You are a senior Python developer. Write a concise code review for the given function. Provide 3-6 actionable points."),
-    review = response.content.strip()
+        ("user", "Here is the function:\n{code}")
+    ])
+    chain = prompt | llm
+    review = chain.invoke({"code": state["code"]})
     return {"draft_review": review}
 
-# ---------- Reflect Node ----------
+# Node: reflect
-# Structured output schema
-SCHEMA = {
-    "pep8": "int (0-10)",
-    "type_hints": "int (0-10)",
-    "edge_cases": "int (0-10)",
-    "naming": "int (0-10)",
-    "weakest_criterion": "string (one of the keys above)",
-    "verdict": "string (\"ok\" or \"needs_revision\")",
-}
 
-parser = StructuredOutputParser.from_function_signature(
+def reflect(state: CodeReviewState) -> Dict[str, object]:
-    "def scores(pep8: int, type_hints: int, edge_cases: int, naming: int, weakest_criterion: str, verdict: str) -> dict"
+    prompt = ChatPromptTemplate.from_messages([
-)
+        ("system", """You are a code quality critic. Score the following review on four criteria: PEP8, type hints, edge cases, naming. Return a JSON with integer scores 0-10, the weakest criterion, and verdict \"ok\" or \"needs_revision\".\n""") ,
-
+        ("user", "Review:\n{draft_review}")
-REFLECT_PROMPT = ChatPromptTemplate.from_messages([
+    ])
-    ("system", "You are a code review critic. Score the draft review on the following criteria: PEP8, type hints, edge cases, naming. Provide scores 0-10 and decide if the review is \"ok\" or \"needs_revision\".")
+    parser = JsonOutputParser()
-])
+    chain = prompt | llm | parser
-
+    result = chain.invoke({"draft_review": state["draft_review"]})
-async def reflect(state: CodeReviewState) -> Dict[str, object]:
+    # result is a dict
-    prompt = REFLECT_PROMPT.format_messages(draft_review=state["draft_review"])
-    response = await llm.ainvoke(prompt)
-    # Parse structured output
-    try:
-        parsed = parser.parse(response.content)
-    except Exception as e:
-        # Fallback: simple heuristic
-        parsed = {
-            "pep8": 5,
-            "type_hints": 5,
-            "edge_cases": 5,
-            "naming": 5,
-            "weakest_criterion": "pep8",
-            "verdict": "needs_revision",
-        }
     return {
         "criteria_scores": {
-            "pep8": parsed["pep8"],
+            "pep8": result["pep8"],
-            "type_hints": parsed["type_hints"],
+            "type_hints": result["type_hints"],
-            "edge_cases": parsed["edge_cases"],
+            "edge_cases": result["edge_cases"],
-            "naming": parsed["naming"],
+            "naming": result["naming"],
         },
-        "weakest_criterion": parsed["weakest_criterion"],
+        "weakest_criterion": result["weakest_criterion"],
-        "verdict": parsed["verdict"],
+        "verdict": result["verdict"],
     }
 
-# ---------- Rewrite Node ----------
+# Node: rewrite
-async def rewrite(state: CodeReviewState) -> Dict[str, str]:
+
-    # Find the weakest criterion and add a focused improvement note
+def rewrite(state: CodeReviewState) -> Dict[str, str]:
-    wc = state["weakest_criterion"]
+    # Increment round
-    improvement = f"\n- Improve {wc.replace('_', ' ')}: Provide more detailed guidance on this aspect."
+    state["round"] += 1
-    new_review = state["draft_review"] + improvement
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a senior Python developer. Rewrite the review to improve the section about {weakest_criterion}. Keep other points unchanged."),
+        ("user", "Original review:\n{draft_review}")
+    ])
+    chain = prompt | llm
+    new_review = chain.invoke({"weakest_criterion": state["weakest_criterion"], "draft_review": state["draft_review"]})
     return {"draft_review": new_review}
 
-# ---------- Graph ----------
+# Build the graph
 builder = StateGraph(CodeReviewState)
-
 builder.add_node("draft_review", draft_review)
 builder.add_node("reflect", reflect)
 builder.add_node("rewrite", rewrite)
 
-builder.set_entry_point("draft_review")
-
 builder.add_edge("draft_review", "reflect")
+# Conditional edge after reflect
 builder.add_conditional_edges(
     "reflect",
-    lambda x: "END" if x["verdict"] == "ok" else "rewrite",
+    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
 )
 builder.add_edge("rewrite", "reflect")
 
-# Stop after max_rounds
+builder.set_entry_point("draft_review")
-builder.add_conditional_edges(
+builder.set_finish_point("END")
-    "reflect",
-    lambda x: "END" if x["round"] >= x["max_rounds"] else "rewrite",
-)
 
 graph = builder.compile()
 
-# ---------- Demo ----------
+# Demo
-async def main():
+if __name__ == "__main__":
-    # Example function to review
-    code = """
     def sort_numbers(arr):
         return sorted(arr)
-    """
+
-    init_state: CodeReviewState = {
+    code = inspect.getsource(sort_numbers)
+    initial_state: CodeReviewState = {
         "code": code,
         "draft_review": "",
         "criteria_scores": {},
@@ -144,14 +99,12 @@ async def main():
         "round": 0,
         "max_rounds": 2,
     }
-    result = await graph.ainvoke(init_state)
+    result = graph.invoke(initial_state)
     print("\n--- Draft Review ---")
     print(result["draft_review"])
     print("\n--- Scores ---")
     print(result["criteria_scores"])
     print("\n--- Verdict ---")
     print(result["verdict"])
-
+    print("\n--- Round ---")
-if __name__ == "__main__":
+    print(result["round"])
-    import asyncio
-    asyncio.run(main())