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

main.py

@@ -0,0 +1,129 @@
+import os
+import asyncio
+import json
+from typing import TypedDict
+
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from langgraph.graph import StateGraph, START, END
+
+from deepagents import create_deep_agent, tool
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+
+# 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,
+)
+
+# State definition
+class ReflectState(TypedDict):
+    question: str
+    draft: str
+    critique: str
+    verdict: str  # ok | needs_revision
+    round: int
+    max_rounds: int
+
+# Node: draft_answer
+def draft_answer(state: ReflectState) -> ReflectState:
+    prompt = f"Write a concise answer (5-10 sentences) to the following question:\n\n{state['question']}"
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft"] = response.content.strip()
+    return state
+
+# Node: reflect
+def reflect(state: ReflectState) -> ReflectState:
+    prompt = f"""You are a critic evaluating the following draft answer. Provide a verdict ('ok' or 'needs_revision') and 2-3 specific points of improvement. Do not provide the revised answer. Use JSON format:
+{{
+  "verdict": "ok" | "needs_revision",
+  "critique": "list of points"
+}}
+Draft:
+{state['draft']}"""
+    response = llm.invoke([HumanMessage(content=prompt)])
+    try:
+        data = json.loads(response.content)
+    except Exception:
+        data = {"verdict": "needs_revision", "critique": "Could not parse critique"}
+    state["critique"] = data.get("critique", "")
+    state["verdict"] = data.get("verdict", "needs_revision")
+    return state
+
+# Node: rewrite
+def rewrite(state: ReflectState) -> ReflectState:
+    prompt = f"""You are revising the draft answer based on the following critique. Produce a revised answer (5-10 sentences). Do not include the critique. Use the critique points to improve clarity, specificity, and remove filler. Draft:\n{state['draft']}\nCritique:\n{state['critique']}"""
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft"] = response.content.strip()
+    state["round"] = state.get("round", 0) + 1
+    return state
+
+# Build the graph
+def build_graph() -> StateGraph:
+    graph = StateGraph(ReflectState)
+    graph.add_node("draft_answer", draft_answer)
+    graph.add_node("reflect", reflect)
+    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 "END"),
+        {
+            "ok": END,
+            "rewrite": "rewrite",
+            "END": END,
+        },
+    )
+    graph.add_edge("rewrite", "reflect")
+    return graph
+
+# Tool that runs the graph
+def answer_question_tool(question: str, max_rounds: int = 2) -> str:
+    graph = build_graph()
+    initial_state: ReflectState = {
+        "question": question,
+        "draft": "",
+        "critique": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": max_rounds,
+    }
+    final_state = graph.invoke(initial_state)
+    return final_state["draft"]
+
+# DeepAgent tool
+@tool
+def answer_question(query: str) -> str:
+    """Answer a question using a self-reflective process."""
+    return answer_question_tool(query)
+
+# Backend for DeepAgent
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
+
+# Create the DeepAgent
+agent = create_deep_agent(
+    model=llm,
+    tools=[answer_question],
+    backend=backend,
+    system_prompt="You are an assistant that answers questions using a self-reflective process. Use the tool 'answer_question' to answer the question.",
+)
+
+# CLI demo
+async def main():
+    question = "Объясни студенту разницу между tool и resource в MCP"
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=question)]},
+        {"configurable": {"thread_id": "session-1"}},
+    )
+    print(result["messages"][-1].content)
+
+if __name__ == "__main__":
+    asyncio.run(main())