Решение готово к публикации: update main.py

main.py

@@ -1,14 +1,17 @@
-"""LangGraph code review agent.
+"""LangGraph Code Review Agent
 
-Implementation follows assignment:
+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:
-- State: CodeReviewState with 4 criteria.
+
-- Nodes: draft_review, reflect, rewrite.
+1. PEP8 compliance
-- Graph: START -> draft_review -> reflect
+2. Type hints
-  - ok -> END
+3. Edge case handling
-  - needs_revision & round < max_rounds -> rewrite -> reflect
+4. Naming conventions
-- Uses LangGraph and LangChain OpenAI for LLM calls.
+
-- Structured output for critique via Pydantic model.
+If the critic returns "needs_revision" the rewrite node improves the weakest part of the review. The process repeats up to ``max_rounds`` times.
-- Demo function sort_numbers.
+
+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
@@ -16,9 +19,11 @@ from __future__ import annotations
 from typing import TypedDict, Dict
 
 from langgraph.graph import StateGraph, END
-from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_structured_output_node
+from langchain_ollama import ChatOllama
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StructuredOutputParser
 from langchain_core.messages import HumanMessage
-from pydantic import BaseModel, Field
 
 # ---------- State ----------
 class CodeReviewState(TypedDict):
@@ -26,101 +31,112 @@ class CodeReviewState(TypedDict):
     draft_review: str
     criteria_scores: Dict[str, int]
     weakest_criterion: str
-    verdict: str
+    verdict: str  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
 # ---------- LLM ----------
-llm = ChatOpenAI(temperature=0)
+# Use Ollama; adjust model name if needed
+llm = ChatOllama(model="llama3")
 
-# ---------- Nodes ----------
+# ---------- Draft Review Node ----------
-class DraftReviewOutput(BaseModel):
+DRAFT_PROMPT = ChatPromptTemplate.from_messages([
-    review: str = Field(..., description="Draft review text")
+    ("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.")
+])
 
-class ReflectOutput(BaseModel):
+async def draft_review(state: CodeReviewState) -> Dict[str, str]:
-    scores: Dict[str, int] = Field(..., description="Scores 0-10 for each criterion")
+    prompt = DRAFT_PROMPT.format_messages(code=state["code"])
-    weakest_criterion: str = Field(..., description="Criterion with lowest score")
+    response = await llm.ainvoke(prompt)
-    verdict: str = Field(..., description="'ok' or 'needs_revision'")
+    review = response.content.strip()
+    return {"draft_review": review}
 
+# ---------- Reflect Node ----------
+# 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\")",
+}
 
-def draft_review(state: CodeReviewState) -> CodeReviewState:
+parser = StructuredOutputParser.from_function_signature(
-    code = state["code"]
+    "def scores(pep8: int, type_hints: int, edge_cases: int, naming: int, weakest_criterion: str, verdict: str) -> dict"
-    prompt = f"""
+)
-    Write a concise code review (3-6 bullet points) for the following Python function. Focus on style, correctness, and potential improvements.
 
-    ```python
+REFLECT_PROMPT = ChatPromptTemplate.from_messages([
-    {code}
+    ("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\".")
-    ```
+])
-    """
-    response = llm.invoke([HumanMessage(content=prompt)])
-    data = DraftReviewOutput.model_validate_json(response.content)
-    state["draft_review"] = data.review
-    return state
 
+async def reflect(state: CodeReviewState) -> Dict[str, object]:
+    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"],
+            "type_hints": parsed["type_hints"],
+            "edge_cases": parsed["edge_cases"],
+            "naming": parsed["naming"],
+        },
+        "weakest_criterion": parsed["weakest_criterion"],
+        "verdict": parsed["verdict"],
+    }
 
-def reflect(state: CodeReviewState) -> CodeReviewState:
+# ---------- Rewrite Node ----------
-    review = state["draft_review"]
+async def rewrite(state: CodeReviewState) -> Dict[str, str]:
-    prompt = f"""
+    # Find the weakest criterion and add a focused improvement note
-    Evaluate the following review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10. Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
+    wc = state["weakest_criterion"]
-
+    improvement = f"\n- Improve {wc.replace('_', ' ')}: Provide more detailed guidance on this aspect."
-    Review:
+    new_review = state["draft_review"] + improvement
-    {review}
+    return {"draft_review": new_review}
-    """
-    response = llm.invoke([HumanMessage(content=prompt)])
-    data = ReflectOutput.model_validate_json(response.content)
-    state["criteria_scores"] = data.scores
-    state["weakest_criterion"] = data.weakest_criterion
-    state["verdict"] = data.verdict
-    return state
-
-
-def rewrite(state: CodeReviewState) -> CodeReviewState:
-    crit = state["weakest_criterion"]
-    review = state["draft_review"]
-    prompt = f"""
-    The review below is weak in the {crit} criterion. Rewrite the review to improve that aspect.
-    Original review:
-    {review}
-    """
-    new_review = llm.invoke([HumanMessage(content=prompt)])
-    state["draft_review"] = new_review.content
-    state["round"] += 1
-    return state
 
 # ---------- 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")
-# After draft_review, go to reflect
+
 builder.add_edge("draft_review", "reflect")
-# After reflect, decide
 builder.add_conditional_edges(
     "reflect",
-    lambda state: state["verdict"] == "ok",
+    lambda x: "END" if x["verdict"] == "ok" else "rewrite",
-    {"ok": END, "needs_revision": "rewrite"},
 )
-# After rewrite, go back to reflect if rounds remain
+builder.add_edge("rewrite", "reflect")
+
+# Stop after max_rounds
 builder.add_conditional_edges(
-    "rewrite",
+    "reflect",
-    lambda state: state["round"] < state["max_rounds"],
+    lambda x: "END" if x["round"] >= x["max_rounds"] else "rewrite",
-    {"continue": "reflect", "end": END},
 )
-# Compile graph
 
 graph = builder.compile()
 
 # ---------- Demo ----------
-
+async def main():
+    # Example function to review
+    code = """
     def sort_numbers(arr):
         return sorted(arr)
-
+    """
-if __name__ == "__main__":
+    init_state: CodeReviewState = {
-    code_str = "def sort_numbers(arr):\n    return sorted(arr)\n"
+        "code": code,
-    initial_state: CodeReviewState = {
-        "code": code_str,
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
@@ -128,6 +144,14 @@ if __name__ == "__main__":
         "round": 0,
         "max_rounds": 2,
     }
-    result = graph.invoke(initial_state)
+    result = await graph.ainvoke(init_state)
-    print("Final state:")
+    print("\n--- Draft Review ---")
-    print(result)
+    print(result["draft_review"])
+    print("\n--- Scores ---")
+    print(result["criteria_scores"])
+    print("\n--- Verdict ---")
+    print(result["verdict"])
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())