diff --git a/main.py b/main.py index 2cfe23b..46beaaa 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,14 @@ import os import asyncio from typing import TypedDict, Annotated + +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 langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, AIMessage -from deepagents import create_deep_agent -from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend # ---------- LLM ---------- llm = ChatOpenAI( @@ -16,105 +18,116 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- State ---------- -class ReflectState(TypedDict): - question: str - draft: str - critique: str - verdict: str # ok | needs_revision - 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']}" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state['draft'] = response.content - 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." - ) - 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: - 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 ---------- -builder = StateGraph(ReflectState) -builder.add_node("draft_answer", draft_answer) -builder.add_node("reflect", reflect) -builder.add_node("rewrite", rewrite) - -builder.set_entry_point("draft_answer") - -# Transition logic -def should_rewrite(state: ReflectState) -> str: - if state['verdict'] == "ok": - return "END" - if state['round'] >= state['max_rounds']: - return "END" - return "rewrite" - -builder.add_conditional_edges("reflect", should_rewrite, { - "rewrite": "rewrite", - "END": "END", -}) - -builder.add_edge("rewrite", "reflect") - -graph = builder.compile() - -# ---------- DeepAgent wrapper ---------- +# ---------- Backend for deepagents (not used directly in graph but required by create_deep_agent) ---------- 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.", -) +# ---------- State definition ---------- +class ReflectState(TypedDict): + question: str + draft: str + critique: str + verdict: str # "ok" | "needs_revision" + round: int + max_rounds: int -# ---------- CLI ---------- -async def run_graph(question: str, max_rounds: int = 2): +# ---------- Nodes ---------- +async def draft_answer(state: ReflectState) -> ReflectState: + prompt = f"Write a concise answer (5–10 sentences) to the following question:\n\n{state['question']}" + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state["draft"] = response.content.strip() + 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:\n{state['draft']}\n\nProvide a verdict (ok or needs_revision) and 2–3 bullet points of critique." + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + text = response.content.strip() + # Simple parsing: first line verdict, rest critique + lines = text.splitlines() + verdict_line = lines[0].lower() + verdict = "ok" if "ok" in verdict_line else "needs_revision" + critique = "\n".join(lines[1:]).strip() + 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. Keep the answer concise (5–10 sentences).\n\nCritique:\n{state['critique']}\n\nOriginal Draft:\n{state['draft']}" + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state["draft"] = response.content.strip() + 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) + +# Entry point +graph.set_entry_point("draft_answer") + +# Transition logic +# After draft_answer -> 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": + return "END" + if state["round"] < state["max_rounds"]: + return "rewrite" + return "END" + +graph.add_conditional_edges("reflect", reflect_conditional, { + "rewrite": "rewrite", + "END": "END", +}) + +# After rewrite -> reflect +graph.add_edge("rewrite", "reflect") + +app = graph.compile() + +# ---------- DeepAgent wrapper ---------- +# The deepagent will simply forward the human message to the graph and return the final draft. +@tool +def run_graph(question: str) -> str: + """Run the reflection graph for a given question.""" initial_state: ReflectState = { "question": question, "draft": "", "critique": "", "verdict": "", "round": 0, - "max_rounds": max_rounds, + "max_rounds": 2, } - result = await graph.ainvoke(initial_state) - return result + result = app.invoke(initial_state) + return result["draft"] + +agent = create_deep_agent( + model=llm, + tools=[run_graph], + backend=backend, + system_prompt="You are an assistant that can answer questions using a self‑checking process.", +) async def main(): - 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"]) + question = "Объясни студенту разницу между tool и resource в MCP." + response = await agent.ainvoke( + {"messages": [HumanMessage(content=question)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print("Final answer:\n", response["messages"][-1].content) if __name__ == "__main__": asyncio.run(main())