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

main.py

@@ -1,15 +1,6 @@
-#!/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
-from typing import TypedDict, Annotated
 
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
@@ -17,8 +8,6 @@ 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 langchain_core.output_parsers import PydanticOutputParser
-from pydantic import BaseModel, Field
 
 # LLM configuration - OpenRouter
 llm = ChatOpenAI(
@@ -36,8 +25,7 @@ backend = CompositeBackend(
     ]
 )
 
-# ---------- LangGraph components ----------
+# ---------- LangGraph definition ----------
-
 class ReflectState(TypedDict):
     question: str
     draft: str
@@ -46,61 +34,70 @@ class ReflectState(TypedDict):
     round: int
     max_rounds: int
 
-class CritiqueOutput(BaseModel):
+def draft_answer(state: ReflectState) -> ReflectState:
-    verdict: str = Field(description="ok or needs_revision")
+    prompt = f"Write a short answer (5-10 sentences) to the following question: {state['question']}"
-    critique: str = Field(description="2-3 bullet points of critique")
+    msg = HumanMessage(content=prompt)
-
+    response = llm.invoke([msg])
-critique_parser = PydanticOutputParser(pydantic_object=CritiqueOutput)
+    state["draft"] = response.content
-
+    state["round"] = 0
-def draft_answer(state: dict) -> dict:
-    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:
+def reflect(state: ReflectState) -> ReflectState:
     prompt = (
-        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of fluff.\n\nDraft:\n{state['draft']}\n\n"
+        f"You are a critic. Evaluate the following draft answer:\n{state['draft']}\n\n"
-        "Respond with JSON containing 'verdict' ('ok' or 'needs_revision') and 'critique' (2-3 bullet points)."
+        "Provide verdict 'ok' or 'needs_revision' and 2-3 points of critique."
     )
-    result = llm.invoke([HumanMessage(content=prompt)])
+    msg = HumanMessage(content=prompt)
-    critique = critique_parser.parse(result.content)
+    response = llm.invoke([msg])
-    state["verdict"] = critique.verdict
+    text = response.content.strip()
-    state["critique"] = critique.critique
+    verdict = "ok"
+    critique = ""
+    if "needs_revision" in text.lower():
+        verdict = "needs_revision"
+    # Extract critique after the word 'Critique:' if present
+    lower_text = text.lower()
+    if "critique:" in lower_text:
+        idx = lower_text.find("critique:")
+        critique = text[idx + len("critique:") :].strip()
+    else:
+        parts = text.split("\n")
+        if len(parts) > 1:
+            critique = "\n".join(parts[1:]).strip()
+    state["verdict"] = verdict
+    state["critique"] = critique
     return state
 
-def rewrite(state: dict) -> dict:
+def rewrite(state: ReflectState) -> ReflectState:
     prompt = (
-        f"Rewrite the draft answer to address the following critique:\n\n{state['critique']}\n\n"
+        f"Rewrite the draft answer to address the following critique:\n{state['critique']}\n\n"
-        "Keep the answer brief (5-10 sentences)."
+        "Keep the answer short (5-10 sentences)."
     )
-    result = llm.invoke([HumanMessage(content=prompt)])
+    msg = HumanMessage(content=prompt)
-    state["draft"] = result.content.strip()
+    response = llm.invoke([msg])
+    state["draft"] = response.content
     state["round"] += 1
     return state
 
-def build_graph() -> StateGraph:
+graph = StateGraph(ReflectState)
-    graph = StateGraph(ReflectState)
+graph.add_node("draft_answer", draft_answer)
-    graph.add_node("draft_answer", draft_answer)
+graph.add_node("reflect", reflect)
-    graph.add_node("reflect", reflect)
+graph.add_node("rewrite", rewrite)
-    graph.add_node("rewrite", rewrite)
+graph.set_entry_point("draft_answer")
+graph.add_edge("draft_answer", "reflect")
+graph.add_conditional_edges(
+    "reflect",
+    lambda s: (
+        "ok"
+        if s["verdict"] == "ok"
+        else ("rewrite" if s["round"] < s["max_rounds"] else "stop")
+    ),
+    {"ok": END, "rewrite": "rewrite", "stop": END},
+)
+graph.add_edge("rewrite", "reflect")
+compiled_graph = graph.compile()
 
-    graph.add_edge(START, "draft_answer")
+def run_reflection(question: str) -> str:
-    graph.add_edge("draft_answer", "reflect")
+    state: ReflectState = {
-
-    def decide_next(state: dict) -> str:
-        if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
-            return "rewrite"
-        return END
-
-    graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "END": END})
-    graph.add_edge("rewrite", "reflect")
-
-    return graph.compile()
-
-def run_graph(question: str) -> str:
-    graph = build_graph()
-    initial_state: ReflectState = {
         "question": question,
         "draft": "",
         "critique": "",
@@ -108,44 +105,42 @@ def run_graph(question: str) -> str:
         "round": 0,
         "max_rounds": 2,
     }
-    final_state = graph.run(initial_state)
+    final_state = compiled_graph.invoke(state)
-    return final_state["draft"]
+    output = (
+        f"Draft:\n{final_state['draft']}\n\n"
+        f"Critique:\n{final_state['critique']}\n\n"
+        f"Verdict: {final_state['verdict']}\n\n"
+        f"Final answer:\n{final_state['draft']}\n"
+    )
+    return output
 
 # ---------- DeepAgents tool ----------
-
 @tool
-def run_graph_tool(question: str) -> str:
+def answer_with_reflection(question: str) -> str:
-    """Run the LangGraph to produce a refined answer."""
+    """Generate a short answer with self-reflection and rewrite if needed."""
-    return run_graph(question)
+    return run_reflection(question)
-
-# ---------- Agent ----------
 
+# ---------- DeepAgents agent ----------
 agent = create_deep_agent(
     model=llm,
-    tools=[run_graph_tool],
+    tools=[answer_with_reflection],
     backend=backend,
-    system_prompt="You are a helpful agent that answers questions by running the run_graph tool.",
+    system_prompt=(
+        "You are a helpful agent that writes short answers and self-reflects. "
+        "Use the tool 'answer_with_reflection' to answer questions."
+    ),
 )
 
-# ---------- CLI ----------
+# ---------- Demo ----------
-
 async def main():
-    parser = argparse.ArgumentParser(description="Answer a question with self-reflection.")
+    question = (
-    parser.add_argument("question", nargs="*", help="The question to answer.")
+        "Explain to a student the difference between tool and resource in MCP."
-    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
+    print(result["messages"][-1].content)
-    final_message = result["messages"][-1].content
-    print("\nAnswer:\n")
-    print(final_message)
 
 if __name__ == "__main__":
     asyncio.run(main())