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

main.py

@@ -1,17 +1,26 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# DESIGN DECISION: deepagents is required by the course assignment to build the agent.
+# NECESSITY: The assignment explicitly requires using create_deep_agent from deepagents; without it the agent cannot be instantiated.
+# OPTIMALITY: Using deepagents ensures consistent agent behavior and simplifies tool integration; alternative frameworks would violate the course constraints.
+# ALTERNATIVES CONSIDERED: Using plain langgraph without deepagents would miss the required framework; manually handling tool calls would increase boilerplate.
+
 import os
 import asyncio
+import argparse
 from typing import TypedDict, Annotated
 
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
-from langgraph.graph import StateGraph, START, END
+from langchain.tools import tool
-from langgraph.graph.message import add_messages
 from deepagents import create_deep_agent
-from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langgraph.graph import StateGraph, START, END
+from langchain_core.output_parsers import PydanticOutputParser
+from pydantic import BaseModel, Field
 
-# ----------------------------------------------------------------------
+# LLM configuration - OpenRouter
-# LLM configuration (OpenRouter, required by the course)
-# ----------------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -19,9 +28,7 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ----------------------------------------------------------------------
+# Backend for deepagents
-# Backend for deepagents (required by the framework)
-# ----------------------------------------------------------------------
 backend = CompositeBackend(
     [
         LocalShellBackend(workspace_dir="./workspace"),
@@ -29,134 +36,70 @@ backend = CompositeBackend(
     ]
 )
 
-# ----------------------------------------------------------------------
+# ---------- LangGraph components ----------
-# State definition for the reflection loop
+
-# ----------------------------------------------------------------------
 class ReflectState(TypedDict):
     question: str
     draft: str
     critique: str
-    verdict: str  # "ok" or "needs_revision"
+    verdict: str  # ok | needs_revision
     round: int
     max_rounds: int
 
-# ----------------------------------------------------------------------
+class CritiqueOutput(BaseModel):
-# Helper function to build a simple LLM chain for a given prompt
+    verdict: str = Field(description="ok or needs_revision")
-# ----------------------------------------------------------------------
+    critique: str = Field(description="2-3 bullet points of critique")
-def llm_call(prompt: str, state: ReflectState) -> str:
-    """Invoke the LLM with a system prompt and the current state."""
-    messages = [
-        HumanMessage(content=prompt.format(**state))
-    ]
-    response = llm.invoke(messages)
-    return response.content.strip()
 
-# ----------------------------------------------------------------------
+critique_parser = PydanticOutputParser(pydantic_object=CritiqueOutput)
-# Node: draft_answer - produce the first answer
+
-# ----------------------------------------------------------------------
+def draft_answer(state: dict) -> dict:
-def draft_answer(state: ReflectState) -> ReflectState:
+    prompt = f"Write a brief answer (5-10 sentences) to the following question:\n\n{state['question']}"
+    result = llm.invoke([HumanMessage(content=prompt)])
+    state["draft"] = result.content.strip()
+    return state
+
+def reflect(state: dict) -> dict:
     prompt = (
-        "You are an expert educator. Answer the following question in 5-10 sentences, "
+        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of fluff.\n\nDraft:\n{state['draft']}\n\n"
-        "clear and concise, without unnecessary filler. Question: {question}"
+        "Respond with JSON containing 'verdict' ('ok' or 'needs_revision') and 'critique' (2-3 bullet points)."
     )
-    draft = llm_call(prompt, state)
+    result = llm.invoke([HumanMessage(content=prompt)])
-    state["draft"] = draft
+    critique = critique_parser.parse(result.content)
-    state["round"] = 0
+    state["verdict"] = critique.verdict
+    state["critique"] = critique.critique
     return state
 
-# ----------------------------------------------------------------------
+def rewrite(state: dict) -> dict:
-# Node: reflect - LLM critic evaluates the draft
+    prompt = (
-# ----------------------------------------------------------------------
+        f"Rewrite the draft answer to address the following critique:\n\n{state['critique']}\n\n"
-def reflect(state: ReflectState) -> ReflectState:
+        "Keep the answer brief (5-10 sentences)."
-    critique_prompt = (
-        "You are a reviewer. Evaluate the draft answer provided below. "
-        "Assess completeness, concreteness and absence of filler. "
-        "Return a verdict ('ok' or 'needs_revision') and list 2-3 short remarks. "
-        "Format exactly as:\n"
-        "Verdict: <verdict>\n"
-        "Critique:\n"
-        "- <remark 1>\n"
-        "- <remark 2>\n"
-        "Draft:\n{draft}"
     )
-    critique_raw = llm_call(critique_prompt, state)
+    result = llm.invoke([HumanMessage(content=prompt)])
-    # Parse the structured response
+    state["draft"] = result.content.strip()
-    lines = critique_raw.splitlines()
-    verdict_line = next((l for l in lines if l.lower().startswith("verdict:")), "")
-    verdict = verdict_line.split(":", 1)[1].strip().lower()
-    critique_start = lines.index("Critique:") + 1 if "Critique:" in lines else 0
-    critique_items = [l.lstrip("- ").strip() for l in lines[critique_start:] if l.startswith("-")]
-    state["verdict"] = verdict
-    state["critique"] = "\n".join(critique_items)
-    return state
-
-# ----------------------------------------------------------------------
-# Node: rewrite - improve the draft based on critique
-# ----------------------------------------------------------------------
-def rewrite(state: ReflectState) -> ReflectState:
-    rewrite_prompt = (
-        "You are a writer. Improve the previous draft according to the following critique points. "
-        "Make the answer clearer, more concrete and remove any filler. Keep the length 5-10 sentences.\n"
-        "Critique:\n{critique}\n\nCurrent draft:\n{draft}"
-    )
-    new_draft = llm_call(rewrite_prompt, state)
-    state["draft"] = new_draft
     state["round"] += 1
     return state
 
-# ----------------------------------------------------------------------
-# DESIGN DECISION: Use a pure LangGraph state machine for reflection.
-# NECESSITY: The assignment explicitly requires a separate reflect node and
-#           iteration via rewrite → reflect, not a try/except retry loop.
-# OPTIMALITY: Graph representation makes the flow declarative, guarantees
-#           max_rounds enforcement, and isolates each responsibility.
-# ALTERNATIVES CONSIDERED: A manual while-loop with try/except was removed
-#           because it mixes error handling with logical revision, violating
-#           the task specification.
-# ----------------------------------------------------------------------
 def build_graph() -> StateGraph:
     graph = StateGraph(ReflectState)
-
     graph.add_node("draft_answer", draft_answer)
     graph.add_node("reflect", reflect)
     graph.add_node("rewrite", rewrite)
 
-    # START → draft_answer
     graph.add_edge(START, "draft_answer")
-    # draft_answer → reflect
     graph.add_edge("draft_answer", "reflect")
 
-    # reflect → END if ok
+    def decide_next(state: dict) -> str:
-    graph.add_conditional_edges(
+        if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
-        "reflect",
+            return "rewrite"
-        lambda s: END if s["verdict"] == "ok" else "rewrite",
-    )
-    # rewrite → reflect (if rounds left)
-    def rewrite_condition(s: ReflectState):
-        if s["round"] < s["max_rounds"]:
-            return "reflect"
         return END
 
-    graph.add_edge("rewrite", rewrite_condition)
+    graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "END": END})
+    graph.add_edge("rewrite", "reflect")
 
-    graph.set_entry_point("draft_answer")
+    return graph.compile()
-    return graph
 
-# ----------------------------------------------------------------------
+def run_graph(question: str) -> str:
-# DeepAgent wrapper - required by the course
+    graph = build_graph()
-# ----------------------------------------------------------------------
-agent = create_deep_agent(
-    model=llm,
-    tools=[],  # No external tools needed for this assignment
-    backend=backend,
-    system_prompt="You are a reflective assistant that writes concise answers and improves them based on critique.",
-)
-
-# ----------------------------------------------------------------------
-# Demo execution
-# ----------------------------------------------------------------------
-async def main():
-    question = "Объясни студенту разницу между tool и resource в MCP"
     initial_state: ReflectState = {
         "question": question,
         "draft": "",
@@ -165,18 +108,44 @@ async def main():
         "round": 0,
         "max_rounds": 2,
     }
+    final_state = graph.run(initial_state)
+    return final_state["draft"]
 
-    graph = build_graph()
+# ---------- DeepAgents tool ----------
-    # Run the graph synchronously (LangGraph supports async, but our nodes are sync)
-    final_state = await graph.ainvoke(initial_state, config={"configurable": {"thread_id": "demo-1"}})
 
-    print("=== Final Answer ===")
+@tool
-    print(final_state["draft"])
+def run_graph_tool(question: str) -> str:
-    print("\n=== Verdict ===")
+    """Run the LangGraph to produce a refined answer."""
-    print(final_state["verdict"])
+    return run_graph(question)
-    if final_state["critique"]:
+
-        print("\n=== Critique ===")
+# ---------- Agent ----------
-        print(final_state["critique"])
+
+agent = create_deep_agent(
+    model=llm,
+    tools=[run_graph_tool],
+    backend=backend,
+    system_prompt="You are a helpful agent that answers questions by running the run_graph tool.",
+)
+
+# ---------- CLI ----------
+
+async def main():
+    parser = argparse.ArgumentParser(description="Answer a question with self-reflection.")
+    parser.add_argument("question", nargs="*", help="The question to answer.")
+    args = parser.parse_args()
+    if not args.question:
+        question = input("Enter your question: ").strip()
+    else:
+        question = " ".join(args.question).strip()
+
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=question)]},
+        {"configurable": {"thread_id": "session-1"}},
+    )
+    # The agent will return the final answer in the last message
+    final_message = result["messages"][-1].content
+    print("\nAnswer:\n")
+    print(final_message)
 
 if __name__ == "__main__":
     asyncio.run(main())