Rebuild the task as a single LangGraph code-review reflection graph with structured scores, targeted rewrite, max_rounds guard, and CLI demo.: update main.py

2026-06-18 13:23:17 +00:00
parent 061d1c231e
commit 25cc4c772d
1 changed files with 173 additions and 79 deletions
@@ -1,97 +1,168 @@
-from typing import TypedDict, Dict
-import inspect
+import os
+import textwrap
+from typing import Literal, TypedDict
+
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, START, StateGraph
+from pydantic import BaseModel, Field

-from langgraph.graph import StateGraph, END
-from langchain_ollama import Ollama
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import JsonOutputParser

-# Define the state
 class CodeReviewState(TypedDict):
    code: str
    draft_review: str
-    criteria_scores: Dict[str, int]
+    criteria_scores: dict[str, int]
    weakest_criterion: str
-    verdict: str  # "ok" | "needs_revision"
+    verdict: str
    round: int
    max_rounds: int

-# LLM instance (Ollama)
-llm = Ollama(model="llama3.1")

-# Node: draft_review
+class ReflectionResult(BaseModel):
+    criteria_scores: dict[str, int] = Field(
+        description="Scores from 0 to 10 for pep8, type_hints, edge_cases, naming."
+    )
+    weakest_criterion: Literal["pep8", "type_hints", "edge_cases", "naming"]
+    verdict: Literal["ok", "needs_revision"]

-def draft_review(state: CodeReviewState) -> Dict[str, str]:
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are a senior Python developer. Write a concise code review for the given function. Provide 3-6 actionable points."),
-        ("user", "Here is the function:\n{code}")
-    ])
-    chain = prompt | llm
-    review = chain.invoke({"code": state["code"]})
-    return {"draft_review": review}

-# Node: reflect
+def build_llm() -> ChatOpenAI:
+    return ChatOpenAI(
+        model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b"),
+        base_url=os.getenv("OPENAI_BASE_URL"),
+        api_key=os.getenv("OPENAI_API_KEY", "dummy"),
+        temperature=0,
+    )

-def reflect(state: CodeReviewState) -> Dict[str, object]:
-    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}")
-    ])
-    parser = JsonOutputParser()
-    chain = prompt | llm | parser
-    result = chain.invoke({"draft_review": state["draft_review"]})
-    # result is a dict
+
+def draft_review(state: CodeReviewState) -> dict:
+    llm = build_llm()
+    prompt = textwrap.dedent(
+        """
+        Ты опытный Python code reviewer.
+        Напиши code review по функции ниже.
+        Требования:
+        - 3-6 конкретных пунктов;
+        - оцени сильные стороны и что улучшить;
+        - обязательно затронь PEP8, type hints, edge cases и naming, если это уместно;
+        - ответ пиши на русском.
+
+        Код:
+        {code}
+        """
+    ).strip()
+    response = llm.invoke(prompt.format(code=state["code"]))
+    return {"draft_review": response.content.strip()}
+
+
+def reflect(state: CodeReviewState) -> dict:
+    llm = build_llm().with_structured_output(ReflectionResult)
+    prompt = textwrap.dedent(
+        """
+        Ты критик качества code review.
+        Оцени review по 4 критериям:
+        - pep8
+        - type_hints
+        - edge_cases
+        - naming
+
+        Правила:
+        - для каждого критерия выставь integer score от 0 до 10;
+        - weakest_criterion — самый слабый критерий;
+        - verdict = "ok", если review уже достаточно хорошее;
+        - verdict = "needs_revision", если самое слабое место стоит усилить.
+
+        Код:
+        {code}
+
+        Review:
+        {review}
+        """
+    ).strip()
+    result = llm.invoke(
+        prompt.format(code=state["code"], review=state["draft_review"])
+    )
    return {
-        "criteria_scores": {
-            "pep8": result["pep8"],
-            "type_hints": result["type_hints"],
-            "edge_cases": result["edge_cases"],
-            "naming": result["naming"],
-        },
-        "weakest_criterion": result["weakest_criterion"],
-        "verdict": result["verdict"],
+        "criteria_scores": result.criteria_scores,
+        "weakest_criterion": result.weakest_criterion,
+        "verdict": result.verdict,
    }

-# Node: rewrite

-def rewrite(state: CodeReviewState) -> Dict[str, str]:
-    # Increment round
-    state["round"] += 1
-    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}
+def rewrite(state: CodeReviewState) -> dict:
+    llm = build_llm()
+    prompt = textwrap.dedent(
+        """
+        Ты улучшаешь уже написанный code review.
+        Нужно целенаправленно усилить самое слабое место: {weakest_criterion}.

-# Build the graph
-builder = StateGraph(CodeReviewState)
-builder.add_node("draft_review", draft_review)
-builder.add_node("reflect", reflect)
-builder.add_node("rewrite", rewrite)
+        Требования:
+        - сохрани общий формат краткого review;
+        - сделай акцент именно на критерии {weakest_criterion};
+        - добавь более точные и полезные замечания;
+        - ответ пиши на русском;
+        - итог должен остаться в формате 3-6 пунктов.

-builder.add_edge("draft_review", "reflect")
-# Conditional edge after reflect
-builder.add_conditional_edges(
-    "reflect",
-    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
-)
-builder.add_edge("rewrite", "reflect")
+        Код:
+        {code}

-builder.set_entry_point("draft_review")
-builder.set_finish_point("END")
+        Текущий review:
+        {review}
+        """
+    ).strip()
+    response = llm.invoke(
+        prompt.format(
+            weakest_criterion=state["weakest_criterion"],
+            code=state["code"],
+            review=state["draft_review"],
+        )
+    )
+    return {
+        "draft_review": response.content.strip(),
+        "round": state["round"] + 1,
+    }

-graph = builder.compile()

-# Demo
-if __name__ == "__main__":
-    def sort_numbers(arr):
-        return sorted(arr)
+def next_step(state: CodeReviewState) -> str:
+    if state["verdict"] == "ok":
+        return "finish"
+    if state["round"] >= state["max_rounds"]:
+        return "finish"
+    return "rewrite"

-    code = inspect.getsource(sort_numbers)
+
+def build_graph():
+    builder = StateGraph(CodeReviewState)
+    builder.add_node("draft_review", draft_review)
+    builder.add_node("reflect", reflect)
+    builder.add_node("rewrite", rewrite)
+    builder.add_edge(START, "draft_review")
+    builder.add_edge("draft_review", "reflect")
+    builder.add_conditional_edges(
+        "reflect",
+        next_step,
+        {
+            "rewrite": "rewrite",
+            "finish": END,
+        },
+    )
+    builder.add_edge("rewrite", "reflect")
+    return builder.compile()
+
+
+def demo_code() -> str:
+    return textwrap.dedent(
+        """
+        def sort_numbers(arr):
+            return sorted(arr)
+        """
+    ).strip()
+
+
+def run_demo() -> None:
+    graph = build_graph()
    initial_state: CodeReviewState = {
-        "code": code,
+        "code": demo_code(),
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
@@ -99,12 +170,35 @@ if __name__ == "__main__":
        "round": 0,
        "max_rounds": 2,
    }
-    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 ---")
-    print(result["round"])
+
+    print("=== Code Review Reflection Demo ===")
+    print("Code under review:")
+    print(initial_state["code"])
+
+    final_state = initial_state.copy()
+    for chunk in graph.stream(initial_state, stream_mode="updates"):
+        for node_name, update in chunk.items():
+            final_state.update(update)
+            print()
+            if node_name == "draft_review":
+                print("--- Draft Review ---")
+                print(final_state["draft_review"])
+            elif node_name == "reflect":
+                print("--- Critic Scores ---")
+                print(final_state["criteria_scores"])
+                print(f"Weakest criterion: {final_state['weakest_criterion']}")
+                print(f"Verdict: {final_state['verdict']}")
+            elif node_name == "rewrite":
+                print(f"--- Rewritten Review After Round {final_state['round']} ---")
+                print(final_state["draft_review"])
+
+    print()
+    print("=== Final Result ===")
+    print(final_state["draft_review"])
+    print(final_state["criteria_scores"])
+    print(f"Rounds used: {final_state['round']} / {final_state['max_rounds']}")
+
+
+if __name__ == "__main__":
+    load_dotenv()
+    run_demo()