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

main.py

@@ -1,38 +1,18 @@
 import os
 import asyncio
-from typing import TypedDict, Dict
+from typing import TypedDict, Annotated, 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
-
-# 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",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
-)
-
-# Backend for deepagents - simple filesystem
-backend = FilesystemBackend()
-
-# 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):
@@ -46,117 +26,127 @@ class CodeReviewState(TypedDict):
 
 # ---------- Structured output for critic ----------
 class CriticOutput(BaseModel):
-    scores: Dict[str, int] = Field(
+    pep8: int = Field(description="Score for PEP8 compliance (0-10)")
-        description="Scores for each criterion: pep8, type_hints, edge_cases, naming. Values 0-10."
+    type_hints: int = Field(description="Score for type hints (0-10)")
-    )
+    edge_cases: int = Field(description="Score for edge case handling (0-10)")
-    verdict: str = Field(
+    naming: int = Field(description="Score for naming conventions (0-10)")
-        description='Verdict: "ok" if all scores >= 7, otherwise "needs_revision".'
+    verdict: str = Field(description='Verdict: "ok" or "needs_revision"')
-    )
 
 critic_parser = PydanticOutputParser(pydantic_object=CriticOutput)
 
-# ---------- Graph nodes ----------
+# ---------- LLM and backend ----------
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
+)
+
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
+
+agent = create_deep_agent(
+    model=llm,
+    tools=[],
+    backend=backend,
+    system_prompt="You are a helpful agent.",
+)
+
+# ---------- Node functions ----------
 async def draft_review(state: CodeReviewState) -> CodeReviewState:
-    prompt = (
+    prompt = f"Please provide a concise code review (3-6 points) for the following Python function:\n\n{state['code']}"
-        f"Write a concise code review (3-6 bullet points) for the following Python function:\n\n"
+    result = await agent.ainvoke([HumanMessage(content=prompt)])
-        f"{state['code']}\n\n"
+    review = result["messages"][-1].content.strip()
-        "Focus on style, correctness, and potential improvements."
+    state["draft_review"] = review
-    )
-    response = await agent.ainvoke(
-        {"messages": [HumanMessage(content=prompt)]},
-        {"configurable": {"thread_id": "draft_review"}},
-    )
-    review_text = response["messages"][-1].content.strip()
-    state["draft_review"] = review_text
     print("\n--- Draft Review ---")
-    print(review_text)
+    print(review)
     return state
 
 async def reflect(state: CodeReviewState) -> CodeReviewState:
     prompt = (
-        f"Evaluate the following code review and assign scores (0-10) for each criterion:\n\n"
+        f"Evaluate the following draft review:\n\n{state['draft_review']}\n\n"
-        f"Review:\n{state['draft_review']}\n\n"
+        "Score each of the following criteria on a scale of 0-10:\n"
-        "Criteria:\n"
+        "- pep8\n- type_hints\n- edge_cases\n- naming\n\n"
-        "1. pep8: adherence to PEP8 style guide.\n"
+        "Return the scores and a verdict ('ok' or 'needs_revision') in the following JSON format:\n"
-        "2. type_hints: presence and correctness of type hints.\n"
+        "{\n  \"pep8\": int,\n  \"type_hints\": int,\n  \"edge_cases\": int,\n  \"naming\": int,\n  \"verdict\": \"ok\" | \"needs_revision\"\n}"
-        "3. edge_cases: handling of edge cases and robustness.\n"
-        "4. naming: clarity and consistency of names.\n\n"
-        "Return a JSON object with keys 'scores' (dict) and 'verdict' ('ok' or 'needs_revision')."
     )
-    response = await agent.ainvoke(
+    result = await agent.ainvoke([HumanMessage(content=prompt)])
-        {"messages": [HumanMessage(content=prompt)]},
+    raw_output = result["messages"][-1].content.strip()
-        {"configurable": {"thread_id": "reflect"}},
-    )
-    raw_output = response["messages"][-1].content.strip()
     try:
         parsed = critic_parser.parse(raw_output)
     except Exception as e:
-        # Fallback: simple parsing if JSON is malformed
+        # Fallback: simple parsing if LLM output is not perfectly formatted
-        import json
+        parsed = CriticOutput(
-        parsed = CriticOutput(**json.loads(raw_output))
+            pep8=0,
-    state["criteria_scores"] = parsed.scores
+            type_hints=0,
-    # Determine weakest criterion
+            edge_cases=0,
-    weakest = min(parsed.scores.items(), key=lambda kv: kv[1])[0]
+            naming=0,
+            verdict="needs_revision",
+        )
+    scores = {
+        "pep8": parsed.pep8,
+        "type_hints": parsed.type_hints,
+        "edge_cases": parsed.edge_cases,
+        "naming": parsed.naming,
+    }
+    weakest = min(scores, key=scores.get)
+    state["criteria_scores"] = scores
     state["weakest_criterion"] = weakest
     state["verdict"] = parsed.verdict
-    print("\n--- Critic Scores ---")
+    print("\n--- Reflection ---")
-    for crit, score in parsed.scores.items():
+    print(f"Scores: {scores}")
-        print(f"{crit}: {score}")
     print(f"Weakest criterion: {weakest}")
     print(f"Verdict: {parsed.verdict}")
     return state
 
 async def rewrite(state: CodeReviewState) -> CodeReviewState:
-    state["round"] += 1
     prompt = (
-        f"Rewrite the part of the review that addresses the weakest criterion "
+        f"Rewrite the section of the draft review that addresses the weakest criterion "
-        f"('{state['weakest_criterion']}') to improve it. Keep the rest of the review unchanged.\n\n"
+        f"('{state['weakest_criterion']}') to improve it. Keep all other parts unchanged.\n\n"
-        f"Original Review:\n{state['draft_review']}\n\n"
+        f"Original draft review:\n\n{state['draft_review']}"
-        "Provide only the updated review."
     )
-    response = await agent.ainvoke(
+    result = await agent.ainvoke([HumanMessage(content=prompt)])
-        {"messages": [HumanMessage(content=prompt)]},
+    new_review = result["messages"][-1].content.strip()
-        {"configurable": {"thread_id": "rewrite"}},
-    )
-    new_review = response["messages"][-1].content.strip()
     state["draft_review"] = new_review
-    print("\n--- Rewritten Review (Round {}) ---".format(state["round"]))
+    state["round"] += 1
+    print("\n--- Rewritten Review ---")
     print(new_review)
     return state
 
-# ---------- Graph construction ----------
+# ---------- Graph ----------
-builder = StateGraph(CodeReviewState)
+def build_graph() -> StateGraph:
+    graph = StateGraph(CodeReviewState)
+    graph.add_node("draft_review", draft_review)
+    graph.add_node("reflect", reflect)
+    graph.add_node("rewrite", rewrite)
 
-builder.add_node("draft_review", draft_review)
+    graph.add_edge(START, "draft_review")
-builder.add_node("reflect", reflect)
+    graph.add_edge("draft_review", "reflect")
-builder.add_node("rewrite", rewrite)
 
-builder.add_edge(START, "draft_review")
+    def reflect_cond(state: CodeReviewState):
-builder.add_edge("draft_review", "reflect")
-
-# Conditional edges after reflect
-def reflect_conditional(state: CodeReviewState):
         if state["verdict"] == "ok":
             return END
         if state["round"] < state["max_rounds"]:
             return "rewrite"
         return END
 
-builder.add_conditional_edges("reflect", reflect_conditional)
+    graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", END: END})
+    graph.add_edge("rewrite", "reflect")
 
-builder.add_edge("rewrite", "reflect")
+    return graph
-
-graph = builder.compile()
 
 # ---------- Demo ----------
 async def main():
     # Sample function to review
-    sample_code = """
+    code_str = """def sort_numbers(arr):
-def sort_numbers(arr):
+    return sorted(arr)"""
-    return sorted(arr)
+
-"""
     initial_state: CodeReviewState = {
-        "code": sample_code.strip(),
+        "code": code_str,
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
@@ -165,12 +155,15 @@ def sort_numbers(arr):
         "max_rounds": 2,
     }
 
-    final_state = await graph.ainvoke(initial_state)
+    graph = build_graph()
+    app = graph.compile()
+    final_state = await app.ainvoke(initial_state)
+
     print("\n=== Final State ===")
+    print(f"Round: {final_state['round']}")
     print(f"Verdict: {final_state['verdict']}")
-    print(f"Rounds performed: {final_state['round']}")
+    print(f"Draft Review:\n{final_state['draft_review']}")
-    print("\nFinal Review:")
+    print(f"Scores: {final_state['criteria_scores']}")
-    print(final_state["draft_review"])
 
 if __name__ == "__main__":
     asyncio.run(main())