fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

main.py

@@ -1,37 +1,15 @@
-"""
-# main.py
-# LangGraph code review agent with reflection and rewrite loop.
-# Implements the task specification without any deepagents dependency.
-# Uses OpenRouter via langchain-openai.
-
 import os
 import asyncio
-from typing import TypedDict, Annotated, Dict
+from typing import TypedDict, Annotated
+
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
-from dotenv import load_dotenv
 
-load_dotenv()
+# ---------- LLM ----------
-
-# ---------------------------------------------------------------------------
-# State definition
-# ---------------------------------------------------------------------------
-class CodeReviewState(TypedDict):
-    code: str
-    draft_review: str
-    criteria_scores: Dict[str, int]
-    weakest_criterion: str
-    verdict: str  # "ok" | "needs_revision"
-    round: int
-    max_rounds: int
-
-# ---------------------------------------------------------------------------
-# LLM setup
-# ---------------------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -39,134 +17,128 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------------------------------------------------------------------------
+# ---------- State ----------
-# Structured output models for reflect node
+class CodeReviewState(TypedDict):
-# ---------------------------------------------------------------------------
+    code: str
-class ReflectOutput(BaseModel):
+    draft_review: str
-    pep8: int = Field(..., ge=0, le=10)
+    criteria_scores: dict[str, int]
-    type_hints: int = Field(..., ge=0, le=10)
+    weakest_criterion: str
-    edge_cases: int = Field(..., ge=0, le=10)
+    verdict: str  # "ok" | "needs_revision"
-    naming: int = Field(..., ge=0, le=10)
+    round: int
-    weakest_criterion: str = Field(...)
+    max_rounds: int
-    verdict: str = Field(..., regex="^(ok|needs_revision)$")
 
-reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+# ---------- Structured output for reflect ----------
+class ReflectionOutput(BaseModel):
+    pep8: int = Field(description="Score 0-10 for PEP8 compliance")
+    type_hints: int = Field(description="Score 0-10 for type hints usage")
+    edge_cases: int = Field(description="Score 0-10 for handling edge cases")
+    naming: int = Field(description="Score 0-10 for naming conventions")
+    weakest_criterion: str = Field(description="Criterion with lowest score")
+    verdict: str = Field(description="'ok' or 'needs_revision'")
 
-# ---------------------------------------------------------------------------
+parser = PydanticOutputParser(pydantic_object=ReflectionOutput)
-# Node implementations
-# ---------------------------------------------------------------------------
-async def draft_review_node(state: CodeReviewState) -> CodeReviewState:
-    """Generate an initial code review with 3–6 bullet points."""
-    prompt = f"""
-You are a senior Python developer. You will write a concise code review for the following function. Provide 3 to 6 bullet points, each starting with a dash.
 
-Function code:
+# ---------- Nodes ----------
-{state['code']}
 
-Review:"""
+def draft_review_node(state: CodeReviewState) -> CodeReviewState:
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    code = state["code"]
-    state['draft_review'] = response.content.strip()
+    prompt = (
-    return state
+        "You are a senior Python developer.\n"
-
+        "Given the following function, write a concise code review (3-6 bullet points).\n"
-async def reflect_node(state: CodeReviewState) -> CodeReviewState:
+        "Focus on style, correctness, edge cases, and naming.\n"
-    """Critic evaluates the draft review on 4 criteria and returns structured scores."""
+        f"Function:\n{code}\n\nReview:"  # LLM will output review
-    prompt = f"""
-You are a code review critic. Evaluate the following draft review on the four criteria below, assigning a score from 0 (worst) to 10 (excellent). Return the scores and the weakest criterion in a JSON format matching the schema:
-
-{reflect_parser.get_format_instructions()}
-
-Draft review:
-{state['draft_review']}
-
-Scores:"""
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
-    try:
-        parsed = reflect_parser.parse(response.content)
-    except Exception as e:
-        # Fallback: treat as all zeros
-        parsed = ReflectOutput(pep8=0, type_hints=0, edge_cases=0, naming=0, 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_node(state: CodeReviewState) -> CodeReviewState:
-    """Rewrite the part of the review that addresses the weakest criterion."""
-    prompt = f"""
-You are a senior Python developer. The following code review has been identified as weak in the criterion: {state['weakest_criterion']}. Rewrite only the section of the review that addresses this criterion, improving clarity and depth. Keep the rest of the review unchanged.
-
-Original review:
-{state['draft_review']}
-
-Rewritten review:"""
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
-    # Replace only the weak section. For simplicity, we replace the whole review.
-    state['draft_review'] = response.content.strip()
-    state['round'] += 1
-    return state
-
-# ---------------------------------------------------------------------------
-# Graph construction
-# ---------------------------------------------------------------------------
-def create_graph() -> StateGraph:
-    graph = StateGraph(CodeReviewState)
-    graph.add_node("draft_review", draft_review_node)
-    graph.add_node("reflect", reflect_node)
-    graph.add_node("rewrite", rewrite_node)
-
-    # Entry point
-    graph.set_entry_point("draft_review")
-
-    # Transitions
-    graph.add_edge("draft_review", "reflect")
-    graph.add_conditional_edges(
-        "reflect",
-        lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END",
     )
-    graph.add_edge("rewrite", "reflect")
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft_review"] = response.content
+    return state
 
-    return graph
 
-# ---------------------------------------------------------------------------
+def reflect_node(state: CodeReviewState) -> CodeReviewState:
-# CLI helper
+    review = state["draft_review"]
-# ---------------------------------------------------------------------------
+    code = state["code"]
-async def run_review(code: str, max_rounds: int = 2) -> CodeReviewState:
+    prompt = (
+        "You are an automated code review critic.\n"
+        "Given the code and its review, assign a score 0-10 for each of the following criteria:\n"
+        "- pep8: PEP8 compliance\n"
+        "- type_hints: use of type hints\n"
+        "- edge_cases: handling of edge cases\n"
+        "- naming: clarity of names\n"
+        "Return the scores, the weakest criterion, and a verdict ('ok' if all scores >=7, else 'needs_revision').\n"
+        f"Code:\n{code}\n\nReview:\n{review}\n\nOutput in JSON with fields: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict."
+    )
+    response = llm.invoke([HumanMessage(content=prompt)])
+    try:
+        out = parser.parse(response.content)
+    except Exception as e:
+        # Fallback: simple parsing if JSON is not strict
+        import json
+        out = json.loads(response.content)
+    state["criteria_scores"] = {
+        "pep8": out.pep8,
+        "type_hints": out.type_hints,
+        "edge_cases": out.edge_cases,
+        "naming": out.naming,
+    }
+    state["weakest_criterion"] = out.weakest_criterion
+    state["verdict"] = out.verdict
+    return state
+
+
+def rewrite_node(state: CodeReviewState) -> CodeReviewState:
+    weakest = state["weakest_criterion"]
+    review = state["draft_review"]
+    code = state["code"]
+    prompt = (
+        "You are a senior Python developer tasked with improving a code review.\n"
+        f"The current review is:\n{review}\n\nThe weakest criterion is '{weakest}'.\n"
+        "Rewrite only the part of the review that addresses this criterion, making it stronger and more specific.\n"
+        "Keep the rest of the review unchanged.\n"
+        "Output only the updated review."
+    )
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft_review"] = response.content
+    state["round"] = state.get("round", 0) + 1
+    return state
+
+# ---------- Graph ----------
+builder = StateGraph(CodeReviewState)
+builder.add_node("draft_review", draft_review_node)
+builder.add_node("reflect", reflect_node)
+builder.add_node("rewrite", rewrite_node)
+
+builder.set_entry_point("draft_review")
+builder.add_edge("draft_review", "reflect")
+builder.add_conditional_edges(
+    "reflect",
+    lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END",
+)
+builder.add_edge("rewrite", "reflect")
+builder.add_edge("END", END)
+
+graph = builder.compile()
+
+# ---------- Demo ----------
+async def main():
+    demo_code = """
+    def sort_numbers(arr):
+        return sorted(arr)
+    """
     initial_state: CodeReviewState = {
-        "code": code,
+        "code": demo_code.strip(),
-        "draft_review": "",
+        "draft_review": "",  # will be filled
         "criteria_scores": {},
         "weakest_criterion": "",
         "verdict": "",
         "round": 0,
-        "max_rounds": max_rounds,
+        "max_rounds": 2,
     }
-    graph = create_graph()
+    result = await graph.ainvoke(initial_state)
-    final_state = await graph.astate(initial_state)
+    print("\n--- Final Review ---")
-    return final_state
-
-# ---------------------------------------------------------------------------
-# Demo main
-# ---------------------------------------------------------------------------
-if __name__ == "__main__":
-    sample_code = """
-def sort_numbers(arr):
-    return sorted(arr)
-"""
-    result = asyncio.run(run_review(sample_code))
-    print("\n=== Initial Draft Review ===")
     print(result["draft_review"])
-    print("\n=== Scores ===")
+    print("\n--- Scores ---")
-    print(result["criteria_scores"])
+    for k, v in result["criteria_scores"].items():
-    print("\n=== Verdict ===")
+        print(f"{k}: {v}")
-    print(result["verdict"])
+    print(f"Verdict: {result['verdict']}")
-    if result["verdict"] == "needs_revision":
+
-        print("\n=== Rewritten Review ===")
+if __name__ == "__main__":
-        print(result["draft_review"])  # after last rewrite
+    asyncio.run(main())
-        print("\n=== Updated Scores ===")
-        print(result["criteria_scores"])
-""