diff --git a/main.py b/main.py index 28bc1ff..2cfe23b 100644 --- a/main.py +++ b/main.py @@ -1,10 +1,10 @@ import os import asyncio from typing import TypedDict, Annotated - +from langgraph.graph import StateGraph, START, END +from langgraph.graph.message import add_messages from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage -from langchain.tools import tool +from langchain_core.messages import HumanMessage, AIMessage from deepagents import create_deep_agent from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend @@ -16,12 +16,6 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Backend ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - # ---------- State ---------- class ReflectState(TypedDict): question: str @@ -32,90 +26,75 @@ class ReflectState(TypedDict): max_rounds: int # ---------- Nodes ---------- -from langgraph.graph import StateGraph, START, END -from langgraph.graph.message import add_messages - -# Helper to format the prompt for each node -DRAFT_PROMPT = """Write a concise answer (5–10 sentences) to the following question: - -{question} -""" - -REFLECT_PROMPT = """You are a critic. Given the draft answer below, evaluate its completeness, specificity, and absence of filler. Respond with: -1. verdict: either "ok" or "needs_revision" -2. critique: 2–3 bullet points explaining what to improve (if any) - -Draft: -{draft} -""" - -REWRITE_PROMPT = """You are revising the draft answer based on the critique. Produce a new draft that addresses the points. Keep the answer concise (5–10 sentences). - -Critique: -{critique} - -Previous draft: -{draft} -""" - -# Node functions async def draft_answer(state: ReflectState) -> ReflectState: - response = await llm.ainvoke([HumanMessage(content=DRAFT_PROMPT.format(question=state["question"]))]) - state["draft"] = response.content.strip() + prompt = f"Write a concise answer (5–10 sentences) to the following question: {state['question']}" + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft'] = response.content return state async def reflect(state: ReflectState) -> ReflectState: - response = await llm.ainvoke([HumanMessage(content=REFLECT_PROMPT.format(draft=state["draft"]))]) - # Parse verdict and critique - text = response.content.strip() - verdict_line = next((l for l in text.splitlines() if l.lower().startswith("verdict:")), "") - critique_lines = [l for l in text.splitlines() if l.startswith("-") or l.startswith("•")] - verdict = verdict_line.split(":",1)[1].strip().lower() if verdict_line else "needs_revision" - critique = "\n".join(critique_lines) if critique_lines else "" - state["verdict"] = verdict - state["critique"] = critique + 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." + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + # Simple parsing: first line verdict, rest critique + lines = response.content.strip().splitlines() + verdict_line = lines[0].lower() + verdict = "ok" if "ok" in verdict_line else "needs_revision" + critique = "\n".join(lines[1:]) if len(lines) > 1 else "" + state['verdict'] = verdict + state['critique'] = critique return state async def rewrite(state: ReflectState) -> ReflectState: - response = await llm.ainvoke([HumanMessage(content=REWRITE_PROMPT.format(critique=state["critique"], draft=state["draft"]))]) - state["draft"] = response.content.strip() - state["round"] += 1 + prompt = ( + f"Rewrite the draft answer taking into account the following critique: {state['critique']}\n\nOriginal draft: {state['draft']}" + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state['draft'] = response.content + state['round'] += 1 return state # ---------- Graph ---------- -graph = StateGraph(ReflectState) -graph.add_node("draft_answer", draft_answer) -graph.add_node("reflect", reflect) -graph.add_node("rewrite", rewrite) +builder = StateGraph(ReflectState) +builder.add_node("draft_answer", draft_answer) +builder.add_node("reflect", reflect) +builder.add_node("rewrite", rewrite) -# Entry point -graph.set_entry_point("draft_answer") +builder.set_entry_point("draft_answer") -# Transitions -# After draft -> reflect -graph.add_edge("draft_answer", "reflect") -# After reflect -# if ok -> END -# if needs_revision and round < max_rounds -> rewrite -# else -> END - -def reflect_conditional(state: ReflectState): - if state["verdict"] == "ok": +# Transition logic +def should_rewrite(state: ReflectState) -> str: + if state['verdict'] == "ok": return "END" - if state["round"] < state["max_rounds"]: - return "rewrite" - return "END" + if state['round'] >= state['max_rounds']: + return "END" + return "rewrite" -graph.add_conditional_edges("reflect", reflect_conditional, {"rewrite": "rewrite", "END": "END"}) -# After rewrite -> reflect -graph.add_edge("rewrite", "reflect") +builder.add_conditional_edges("reflect", should_rewrite, { + "rewrite": "rewrite", + "END": "END", +}) -graph.compile() +builder.add_edge("rewrite", "reflect") + +graph = builder.compile() # ---------- DeepAgent wrapper ---------- -@tool -def run_reflect_graph(question: str, max_rounds: int = 2) -> str: - """Run the reflection graph and return the final draft.""" +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +agent = create_deep_agent( + model=llm, + tools=[], + backend=backend, + system_prompt="You are a helper that runs a reflection graph.", +) + +# ---------- CLI ---------- +async def run_graph(question: str, max_rounds: int = 2): initial_state: ReflectState = { "question": question, "draft": "", @@ -124,23 +103,18 @@ def run_reflect_graph(question: str, max_rounds: int = 2) -> str: "round": 0, "max_rounds": max_rounds, } - result = graph.invoke(initial_state) - return result["draft"] - -agent = create_deep_agent( - model=llm, - tools=[run_reflect_graph], - backend=backend, - system_prompt="You are an assistant that can answer questions and self‑critique using the provided tool.", -) + result = await graph.ainvoke(initial_state) + return result async def main(): - question = "Объясни студенту разницу между tool и resource в MCP." - response = await agent.ainvoke( - {"messages": [HumanMessage(content=f"Please answer: {question}")]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print("Final answer:\n", response["messages"][-1].content) + question = "Объясни студенту разницу между tool и resource в MCP" + result = await run_graph(question) + print("\n--- Final Draft ---\n") + print(result["draft"]) + print("\n--- Critique ---\n") + print(result["critique"]) + print("\n--- Verdict ---\n") + print(result["verdict"]) if __name__ == "__main__": asyncio.run(main())