fix: main.py — Повторный экзамен #2: Граф с рефлексией на код

main.py

@@ -1,20 +1,21 @@
 import os
 import asyncio
-from typing import TypedDict, Annotated, Dict
+from typing import TypedDict, Dict
 
+from dotenv import load_dotenv
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
-from langchain.tools import tool
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-
 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 deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend
 
-# ---------- LLM ----------
+# Load API key from .env
+load_dotenv()
+
+# LLM configuration - OpenRouter
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -22,124 +23,140 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend ----------
+# Backend for deepagents - simple filesystem
-backend = CompositeBackend([
+backend = FilesystemBackend()
-    LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
-])
 
-# ---------- State ----------
+# Create a deepagents agent that will be used inside the graph nodes
+agent = create_deep_agent(
+    model=llm,
+    tools=[],
+    backend=backend,
+    system_prompt="You are a code review assistant.",
+)
+
+# ---------- State definition ----------
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
     criteria_scores: Dict[str, int]
     weakest_criterion: str
-    verdict: str
+    verdict: str  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
-# ---------- Pydantic model for reflect output ----------
+# ---------- Structured output for critic ----------
-class ReflectOutput(BaseModel):
+class CriticOutput(BaseModel):
-    pep8: int = Field(description="Score for PEP8 compliance (0-10)")
+    scores: Dict[str, int] = Field(
-    type_hints: int = Field(description="Score for type hints (0-10)")
+        description="Scores for each criterion: pep8, type_hints, edge_cases, naming. Values 0-10."
-    edge_cases: int = Field(description="Score for edge case handling (0-10)")
+    )
-    naming: int = Field(description="Score for naming conventions (0-10)")
+    verdict: str = Field(
-    weakest_criterion: str = Field(description="The criterion with the lowest score")
+        description='Verdict: "ok" if all scores >= 7, otherwise "needs_revision".'
-    verdict: str = Field(description="'ok' or 'needs_revision'")
+    )
 
-parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+critic_parser = PydanticOutputParser(pydantic_object=CriticOutput)
 
-# ---------- Nodes ----------
+# ---------- Graph nodes ----------
 async def draft_review(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""
+    prompt = (
-    Write a concise code review for the following Python function. Provide 3-6 bullet points.
+        f"Write a concise code review (3-6 bullet points) for the following Python function:\n\n"
-
+        f"{state['code']}\n\n"
-    Function:
+        "Focus on style, correctness, and potential improvements."
-    ```python
+    )
-    {state['code']}
+    response = await agent.ainvoke(
-    ```
+        {"messages": [HumanMessage(content=prompt)]},
-    """
+        {"configurable": {"thread_id": "draft_review"}},
-    response = await llm.ainvoke(HumanMessage(content=prompt))
+    )
-    state['draft_review'] = response.content.strip()
+    review_text = response["messages"][-1].content.strip()
+    state["draft_review"] = review_text
+    print("\n--- Draft Review ---")
+    print(review_text)
     return state
 
 async def reflect(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""
+    prompt = (
-    Evaluate the following code and draft review. Score each of the four criteria on a scale 0-10:
+        f"Evaluate the following code review and assign scores (0-10) for each criterion:\n\n"
-    - PEP8 compliance
+        f"Review:\n{state['draft_review']}\n\n"
-    - Type hints
+        "Criteria:\n"
-    - Edge case handling
+        "1. pep8: adherence to PEP8 style guide.\n"
-    - Naming conventions
+        "2. type_hints: presence and correctness of type hints.\n"
-
+        "3. edge_cases: handling of edge cases and robustness.\n"
-    Provide the scores, identify the weakest criterion, and give a verdict ('ok' or 'needs_revision').
+        "4. naming: clarity and consistency of names.\n\n"
-
+        "Return a JSON object with keys 'scores' (dict) and 'verdict' ('ok' or 'needs_revision')."
-    Code:
+    )
-    ```python
+    response = await agent.ainvoke(
-    {state['code']}
+        {"messages": [HumanMessage(content=prompt)]},
-    ```
+        {"configurable": {"thread_id": "reflect"}},
-
+    )
-    Draft Review:
+    raw_output = response["messages"][-1].content.strip()
-    ```text
-    {state['draft_review']}
-    ```
-
-    Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.
-    """
-    response = await llm.ainvoke(HumanMessage(content=prompt))
     try:
-        parsed = parser.parse(response.content)
+        parsed = critic_parser.parse(raw_output)
     except Exception as e:
-        # Fallback: set all scores to 0 and verdict to needs_revision
+        # Fallback: simple parsing if JSON is malformed
-        parsed = ReflectOutput(pep8=0, type_hints=0, edge_cases=0, naming=0, weakest_criterion="unknown", verdict="needs_revision")
+        import json
-    state['criteria_scores'] = {
+        parsed = CriticOutput(**json.loads(raw_output))
-        "pep8": parsed.pep8,
+    state["criteria_scores"] = parsed.scores
-        "type_hints": parsed.type_hints,
+    # Determine weakest criterion
-        "edge_cases": parsed.edge_cases,
+    weakest = min(parsed.scores.items(), key=lambda kv: kv[1])[0]
-        "naming": parsed.naming,
+    state["weakest_criterion"] = weakest
-    }
+    state["verdict"] = parsed.verdict
-    state['weakest_criterion'] = parsed.weakest_criterion
+    print("\n--- Critic Scores ---")
-    state['verdict'] = parsed.verdict
+    for crit, score in parsed.scores.items():
+        print(f"{crit}: {score}")
+    print(f"Weakest criterion: {weakest}")
+    print(f"Verdict: {parsed.verdict}")
     return state
 
 async def rewrite(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""
+    state["round"] += 1
-    Rewrite the section of the draft review that addresses the weakest criterion: {state['weakest_criterion']}.
+    prompt = (
-    Keep all other parts of the review unchanged.
+        f"Rewrite the part of the review that addresses the weakest criterion "
-
+        f"('{state['weakest_criterion']}') to improve it. Keep the rest of the review unchanged.\n\n"
-    Current Draft Review:
+        f"Original Review:\n{state['draft_review']}\n\n"
-    ```text
+        "Provide only the updated review."
-    {state['draft_review']}
+    )
-    ```
+    response = await agent.ainvoke(
-    """
+        {"messages": [HumanMessage(content=prompt)]},
-    response = await llm.ainvoke(HumanMessage(content=prompt))
+        {"configurable": {"thread_id": "rewrite"}},
-    state['draft_review'] = response.content.strip()
+    )
-    state['round'] += 1
+    new_review = response["messages"][-1].content.strip()
+    state["draft_review"] = new_review
+    print("\n--- Rewritten Review (Round {}) ---".format(state["round"]))
+    print(new_review)
     return state
 
-# ---------- Graph ----------
+# ---------- Graph construction ----------
-review_graph = StateGraph(CodeReviewState)
+builder = StateGraph(CodeReviewState)
-review_graph.add_node("draft_review", draft_review)
-review_graph.add_node("reflect", reflect)
-review_graph.add_node("rewrite", rewrite)
 
-review_graph.set_entry_point("draft_review")
+builder.add_node("draft_review", draft_review)
-review_graph.add_edge("draft_review", "reflect")
+builder.add_node("reflect", reflect)
+builder.add_node("rewrite", rewrite)
+
+builder.add_edge(START, "draft_review")
+builder.add_edge("draft_review", "reflect")
 
 # Conditional edges after reflect
-review_graph.add_conditional_edges(
+def reflect_conditional(state: CodeReviewState):
-    "reflect",
+    if state["verdict"] == "ok":
-    lambda state: "END" if state["verdict"] == "ok" else "rewrite" if state["round"] < state["max_rounds"] else "END",
+        return END
-)
+    if state["round"] < state["max_rounds"]:
-review_graph.add_edge("rewrite", "reflect")
+        return "rewrite"
+    return END
 
-compiled_graph = review_graph.compile()
+builder.add_conditional_edges("reflect", reflect_conditional)
 
-# ---------- Tool ----------
+builder.add_edge("rewrite", "reflect")
-@tool
+
-async def run_code_review(code: str) -> Dict:
+graph = builder.compile()
-    """Run a code review on the provided Python function."""
+
+# ---------- Demo ----------
+async def main():
+    # Sample function to review
+    sample_code = """
+def sort_numbers(arr):
+    return sorted(arr)
+"""
     initial_state: CodeReviewState = {
-        "code": code,
+        "code": sample_code.strip(),
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
@@ -147,37 +164,13 @@ async def run_code_review(code: str) -> Dict:
         "round": 0,
         "max_rounds": 2,
     }
-    final_state = compiled_graph.invoke(initial_state)
-    return {
-        "final_review": final_state["draft_review"],
-        "scores": final_state["criteria_scores"],
-        "weakest_criterion": final_state["weakest_criterion"],
-        "verdict": final_state["verdict"],
-        "rounds": final_state["round"],
-    }
 
-# ---------- Agent ----------
+    final_state = await graph.ainvoke(initial_state)
-agent = create_deep_agent(
+    print("\n=== Final State ===")
-    model=llm,
+    print(f"Verdict: {final_state['verdict']}")
-    tools=[run_code_review],
+    print(f"Rounds performed: {final_state['round']}")
-    backend=backend,
+    print("\nFinal Review:")
-    system_prompt="You are a helpful code review agent.",
+    print(final_state["draft_review"])
-)
-
-# ---------- Demo ----------
-async def main():
-    # Sample function to review
-    code = """
-    def sort_numbers(arr):
-        return sorted(arr)
-    """
-    result = await run_code_review(code)
-    print("\n=== Final Review ===\n")
-    print(result["final_review"])
-    print("\n=== Scores ===\n")
-    for k, v in result["scores"].items():
-        print(f"{k}: {v}")
-    print(f"\nVerdict: {result['verdict']} (Rounds: {result['rounds']})")
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())