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

main.py

@@ -1,20 +1,37 @@
 import os
 import asyncio
 from typing import TypedDict, Annotated, Dict
-
 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
 
-# ---------- State definition ----------
+# Load environment variables
+from dotenv import load_dotenv
+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
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
+
+# State definition
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
@@ -24,129 +41,100 @@ class CodeReviewState(TypedDict):
     round: int
     max_rounds: int
 
-# ---------- Structured output for critic ----------
+# Structured output for reflect node
-class CriticOutput(BaseModel):
+class ReviewScores(BaseModel):
-    pep8: int = Field(description="Score for PEP8 compliance (0-10)")
+    pep8: int = Field(description="Score 0-10 for PEP8 compliance")
-    type_hints: int = Field(description="Score for type hints (0-10)")
+    type_hints: int = Field(description="Score 0-10 for type hints usage")
-    edge_cases: int = Field(description="Score for edge case handling (0-10)")
+    edge_cases: int = Field(description="Score 0-10 for edge case handling")
-    naming: int = Field(description="Score for naming conventions (0-10)")
+    naming: int = Field(description="Score 0-10 for naming conventions")
-    verdict: str = Field(description='Verdict: "ok" or "needs_revision"')
+    verdict: str = Field(description='\"ok\" or \"needs_revision\"')
 
-critic_parser = PydanticOutputParser(pydantic_object=CriticOutput)
+review_parser = PydanticOutputParser(pydantic_object=ReviewScores)
 
-# ---------- LLM and backend ----------
+# Draft review node
-llm = ChatOpenAI(
+def draft_review(state: CodeReviewState) -> CodeReviewState:
-    model="openai/gpt-oss-20b:free",
+    prompt = f"""Write a concise code review for the following Python function. Provide 3-6 bullet points. Do not include any code blocks. The function is:
-    base_url="https://openrouter.ai/api/v1",
+```python
-    api_key=os.getenv("OPENAI_API_KEY"),
+{state["code"]}
-    temperature=0.0,
+```"""
-)
+    response = llm.invoke(prompt)
-
+    state["draft_review"] = response.content.strip()
-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 = f"Please provide a concise code review (3-6 points) for the following Python function:\n\n{state['code']}"
-    result = await agent.ainvoke([HumanMessage(content=prompt)])
-    review = result["messages"][-1].content.strip()
-    state["draft_review"] = review
     print("\n--- Draft Review ---")
-    print(review)
+    print(state["draft_review"])
     return state
 
-async def reflect(state: CodeReviewState) -> CodeReviewState:
+# Reflect node
-    prompt = (
+def reflect(state: CodeReviewState) -> CodeReviewState:
-        f"Evaluate the following draft review:\n\n{state['draft_review']}\n\n"
+    prompt = f"""You are a code review critic. Evaluate the following draft review for the given code. Provide scores 0-10 for each criterion and a verdict. Return JSON matching the schema:
-        "Score each of the following criteria on a scale of 0-10:\n"
+{review_parser.get_format_instructions()}
-        "- pep8\n- type_hints\n- edge_cases\n- naming\n\n"
+Draft review:
-        "Return the scores and a verdict ('ok' or 'needs_revision') in the following JSON format:\n"
+{state["draft_review"]}
-        "{\n  \"pep8\": int,\n  \"type_hints\": int,\n  \"edge_cases\": int,\n  \"naming\": int,\n  \"verdict\": \"ok\" | \"needs_revision\"\n}"
+
-    )
+Code:
-    result = await agent.ainvoke([HumanMessage(content=prompt)])
+```python
-    raw_output = result["messages"][-1].content.strip()
+{state["code"]}
-    try:
+```"""
-        parsed = critic_parser.parse(raw_output)
+    response = llm.invoke(prompt)
-    except Exception as e:
+    parsed = review_parser.parse(response.content)
-        # Fallback: simple parsing if LLM output is not perfectly formatted
+    state["criteria_scores"] = {
-        parsed = CriticOutput(
-            pep8=0,
-            type_hints=0,
-            edge_cases=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
+    state["weakest_criterion"] = min(state["criteria_scores"], key=state["criteria_scores"].get)
     print("\n--- Reflection ---")
-    print(f"Scores: {scores}")
+    print(f"Scores: {state['criteria_scores']}")
-    print(f"Weakest criterion: {weakest}")
+    print(f"Weakest criterion: {state['weakest_criterion']}")
-    print(f"Verdict: {parsed.verdict}")
+    print(f"Verdict: {state['verdict']}")
     return state
 
-async def rewrite(state: CodeReviewState) -> CodeReviewState:
+# Rewrite node
-    prompt = (
+def rewrite(state: CodeReviewState) -> CodeReviewState:
-        f"Rewrite the section of the draft review that addresses the weakest criterion "
+    prompt = f"""Rewrite the section of the draft review that addresses the weakest criterion ({state['weakest_criterion']}) to improve it. Keep other parts unchanged. Return only the updated draft review.
-        f"('{state['weakest_criterion']}') to improve it. Keep all other parts unchanged.\n\n"
+Original draft review:
-        f"Original draft review:\n\n{state['draft_review']}"
+{state['draft_review']}
-    )
+
-    result = await agent.ainvoke([HumanMessage(content=prompt)])
+Code:
-    new_review = result["messages"][-1].content.strip()
+```python
-    state["draft_review"] = new_review
+{state['code']}
+```"""
+    response = llm.invoke(prompt)
+    state["draft_review"] = response.content.strip()
     state["round"] += 1
-    print("\n--- Rewritten Review ---")
+    print("\n--- Rewrite ---")
-    print(new_review)
+    print(state["draft_review"])
     return state
 
-# ---------- Graph ----------
+# Graph definition
 def build_graph() -> StateGraph:
     graph = StateGraph(CodeReviewState)
     graph.add_node("draft_review", draft_review)
     graph.add_node("reflect", reflect)
     graph.add_node("rewrite", rewrite)
 
-    graph.add_edge(START, "draft_review")
+    graph.set_entry_point("draft_review")
     graph.add_edge("draft_review", "reflect")
 
-    def reflect_cond(state: CodeReviewState):
+    def condition(state: CodeReviewState):
         if state["verdict"] == "ok":
-            return END
+            return "END"
-        if state["round"] < state["max_rounds"]:
+        if state["round"] >= state["max_rounds"]:
+            return "END"
         return "rewrite"
-        return END
 
-    graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", END: END})
+    graph.add_conditional_edges("reflect", condition, ["rewrite", "END"])
     graph.add_edge("rewrite", "reflect")
 
     return graph
 
-# ---------- Demo ----------
+# Tool that runs the graph
-async def main():
+@tool
-    # Sample function to review
+def run_review(code: str) -> str:
-    code_str = """def sort_numbers(arr):
+    """Run a code review cycle on the provided Python function."""
-    return sorted(arr)"""
-
     initial_state: CodeReviewState = {
-        "code": code_str,
+        "code": code,
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
@@ -154,16 +142,32 @@ async def main():
         "round": 0,
         "max_rounds": 2,
     }
-
     graph = build_graph()
-    app = graph.compile()
+    final_state = graph.invoke(initial_state)
-    final_state = await app.ainvoke(initial_state)
+    output = f"Final review after {final_state['round']} round(s):\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}\nVerdict: {final_state['verdict']}"
+    return output
 
-    print("\n=== Final State ===")
+# Deepagents agent
-    print(f"Round: {final_state['round']}")
+agent = create_deep_agent(
-    print(f"Verdict: {final_state['verdict']}")
+    model=llm,
-    print(f"Draft Review:\n{final_state['draft_review']}")
+    tools=[run_review],
-    print(f"Scores: {final_state['criteria_scores']}")
+    backend=backend,
+    system_prompt="You are a helpful code review agent.",
+)
+
+# Demo function
+demo_code = """
+def sort_numbers(arr):
+    return sorted(arr)
+"""
+
+async def main():
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=demo_code)]},
+        {"configurable": {"thread_id": "session-1"}},
+    )
+    print("\n=== Final Output ===")
+    print(result["messages"][-1].content)
 
 if __name__ == "__main__":
     asyncio.run(main())