Update main.py

main.py

@@ -1,16 +1,19 @@
+"""
+# main.py
+# LangGraph agent with reflection and rewrite loop
+# Author: ChatGPT
+# Requirements: langgraph, langchain-openai, deepagents
+
 import os
 import asyncio
-from typing import TypedDict, Annotated
+from typing import TypedDict
 
 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 deepagents import create_deep_agent
+from deepagents.backends import CompositeBackend, LocalShellBackend
 
-# ---------- LLM ----------
+# LLM setup
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -18,103 +21,105 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend for deepagents ----------
+backend = CompositeBackend(
-backend = CompositeBackend([
+    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
-    LocalShellBackend(workspace_dir="./workspace"),
+    routes={},
-    FilesystemBackend(),
+)
-])
+agent = create_deep_agent(
+    model=llm,
+    tools=[],
+    backend=backend,
+    system_prompt="You are a helpful assistant.",
+)
 
-# ---------- State definition ----------
 class ReflectState(TypedDict):
     question: str
     draft: str
     critique: str
-    verdict: str  # ok | needs_revision
+    verdict: str
     round: int
     max_rounds: int
 
-# ---------- Nodes ----------
 async def draft_answer(state: ReflectState) -> ReflectState:
-    prompt = f"Write a concise answer (5–10 sentences) to the following question: {state['question']}"
+    prompt = f"Write a concise answer (5–10 sentences) to the following question:\n\n{state['question']}"
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "draft"}})
-    state['draft'] = response.content.strip()
+    draft = response["messages"][-1].content
+    state["draft"] = draft
     return state
 
 async def reflect(state: ReflectState) -> ReflectState:
     prompt = (
-        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n\nDraft: {state['draft']}\n\nProvide a verdict (ok or needs_revision) and 2–3 bullet points of critique."
+        "You are a critical reviewer.\n"
+        "Evaluate the following draft answer for completeness, specificity, and lack of filler.\n"
+        "Provide a verdict: 'ok' if the answer is satisfactory, otherwise 'needs_revision'.\n"
+        "If revision is needed, give 2–3 concrete points for improvement.\n"
+        "Respond in JSON with keys 'verdict' and 'critique'.\n"
+        f"Draft: {state['draft']}"
     )
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "reflect"}})
-    text = response.content.strip()
+    import json
-    # Simple parsing: first line verdict, rest critique
+    try:
-    lines = text.splitlines()
+        data = json.loads(response["messages"][-1].content)
-    verdict_line = lines[0].lower()
+        verdict = data.get("verdict", "needs_revision")
-    verdict = "ok" if "ok" in verdict_line else "needs_revision"
+        critique = data.get("critique", "")
-    critique = "\n".join(lines[1:]).strip()
+    except Exception:
-    state['verdict'] = verdict
+        verdict = "needs_revision"
-    state['critique'] = critique
+        critique = "Could not parse critique."
+    state["verdict"] = verdict
+    state["critique"] = critique
     return state
 
 async def rewrite(state: ReflectState) -> ReflectState:
     prompt = (
-        f"Rewrite the draft answer taking into account the following critique: {state['critique']}\n\nOriginal draft: {state['draft']}"
+        "You are revising the following draft answer based on the critique.\n"
+        "Make the answer clearer, more specific, and remove any filler.\n"
+        "Do not add new information beyond what is already in the draft.\n"
+        f"Draft: {state['draft']}\n"
+        f"Critique: {state['critique']}"
     )
-    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "rewrite"}})
-    state['draft'] = response.content.strip()
+    new_draft = response["messages"][-1].content
-    state['round'] += 1
+    state["draft"] = new_draft
+    state["round"] += 1
     return state
 
-# ---------- Graph ----------
+
-async def run_graph(question: str, max_rounds: int = 2) -> str:
+def build_graph() -> StateGraph[ReflectState]:
     graph = StateGraph(ReflectState)
     graph.add_node("draft_answer", draft_answer)
     graph.add_node("reflect", reflect)
     graph.add_node("rewrite", rewrite)
-
+    graph.add_edge(START, "draft_answer")
-    graph.set_entry_point("draft_answer")
     graph.add_edge("draft_answer", "reflect")
     graph.add_conditional_edges(
         "reflect",
-        lambda x: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END",
+        lambda state: state["verdict"],
+        {"ok": END, "needs_revision": "rewrite"},
     )
-    graph.add_edge("rewrite", "reflect")
+    graph.add_conditional_edges(
+        "rewrite",
+        lambda state: "rewrite" if state["round"] < state["max_rounds"] else END,
+        {"rewrite": "reflect", END: END},
+    )
+    return graph
 
-    graph.add_edge("END", END)
+async def main():
-
+    question = "Объясни студенту разницу между tool и resource в MCP"
-    app = graph.compile()
     initial_state: ReflectState = {
         "question": question,
         "draft": "",
         "critique": "",
         "verdict": "",
         "round": 0,
-        "max_rounds": max_rounds,
+        "max_rounds": 2,
     }
-    final_state = await app.ainvoke(initial_state)
+    graph = build_graph()
-    return final_state["draft"]
+    result = await graph.ainvoke(initial_state)
-
+    print("\n--- Final Answer ---")
-# ---------- DeepAgent tool ----------
+    print(result["draft"])
-@tool
+    print("\n--- Final Critique ---")
-async def answer_question(query: str) -> str:
+    print(result["critique"])
-    """Generate a refined answer using self‑reflection graph."""
-    return await run_graph(query)
-
-# ---------- DeepAgent ----------
-agent = create_deep_agent(
-    model=llm,
-    tools=[answer_question],
-    backend=backend,
-    system_prompt="You are an AI assistant that answers questions. Use the provided tool to generate answers.",
-)
-
-# ---------- CLI ----------
-async def main():
-    question = "Объясни студенту разницу между tool и resource в MCP"
-    result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=question)]},
-        {"configurable": {"thread_id": "session-1"}},
-    )
-    print("\nFinal answer:\n", result["messages"][-1].content)
 
 if __name__ == "__main__":
     asyncio.run(main())
+"""