From 965eff9436c3b7552271a203b8161b26fae72f25 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=BB=D0=B5=D0=B1=20=D0=9D=D0=B8=D0=BA=D0=B8=D1=88?= =?UTF-8?q?=D0=B8=D0=BD?= Date: Fri, 5 Jun 2026 14:32:20 +0000 Subject: [PATCH] Add main.py --- main.py | 143 ++++++++++++++++++++++---------------------------------- 1 file changed, 56 insertions(+), 87 deletions(-) diff --git a/main.py b/main.py index 46beaaa..8f939ae 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,13 @@ import os +import re 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 -# ---------- LLM ---------- +# LLM configuration (OpenRouter) llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -18,91 +15,78 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Backend for deepagents (not used directly in graph but required by create_deep_agent) ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - # ---------- State definition ---------- class ReflectState(TypedDict): question: str draft: str critique: str - verdict: str # "ok" | "needs_revision" + 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:\n\n{state['question']}" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state["draft"] = response.content.strip() - return state - -async def reflect(state: ReflectState) -> ReflectState: +# ---------- Node implementations ---------- +async def draft_answer(state: ReflectState) -> dict: 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." + f"Write a concise answer (5–10 sentences) to the following question:\n\n" + f"Question: {state['question']}" ) - response = await llm.ainvoke([HumanMessage(content=prompt)]) + response = await llm.ainvoke([{"role": "user", "content": prompt}]) + draft = response.content.strip() + return {"draft": draft, "round": 0} + +async def reflect(state: ReflectState) -> dict: + prompt = ( + f"You are a critical reviewer. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n\n" + f"Draft: {state['draft']}\n\n" + f"Provide a verdict (ok or needs_revision) and 2–3 bullet points of critique." + ) + response = await llm.ainvoke([{"role": "user", "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 + verdict_match = re.search(r"(ok|needs_revision)", text, re.IGNORECASE) + verdict = verdict_match.group(1).lower() if verdict_match else "needs_revision" + return {"critique": text, "verdict": verdict} -async def rewrite(state: ReflectState) -> ReflectState: +async def rewrite(state: ReflectState) -> dict: 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']}" + f"Rewrite the draft answer incorporating the following critique. Keep the answer concise (5–10 sentences).\n\n" + f"Critique: {state['critique']}\n\n" + f"Original Draft: {state['draft']}" ) - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state["draft"] = response.content.strip() - state["round"] += 1 - return state + response = await llm.ainvoke([{"role": "user", "content": prompt}]) + new_draft = response.content.strip() + return {"draft": new_draft, "round": state['round'] + 1} -# ---------- Graph ---------- -graph = StateGraph(ReflectState) -graph.add_node("draft_answer", draft_answer) -graph.add_node("reflect", reflect) -graph.add_node("rewrite", rewrite) +# ---------- Graph construction ---------- +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") +builder.add_edge("draft_answer", "reflect") +builder.add_conditional_edges( + "reflect", + lambda x: x["verdict"], + { + "ok": END, + "needs_revision": "rewrite", + }, +) +builder.add_edge("rewrite", "reflect") -# 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 +# Limit rounds +async def limit_rounds(state: ReflectState) -> str: + if state["round"] >= state["max_rounds"] and state["verdict"] == "needs_revision": + return END + return "reflect" -def reflect_conditional(state: ReflectState): - if state["verdict"] == "ok": - return "END" - if state["round"] < state["max_rounds"]: - return "rewrite" - return "END" +builder.add_conditional_edges("rewrite", limit_rounds, {"reflect": "reflect", END: END}) -graph.add_conditional_edges("reflect", reflect_conditional, { - "rewrite": "rewrite", - "END": "END", -}) +graph = builder.compile() -# 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.""" +# ---------- Demo execution ---------- +async def main(): + question = "Объясни студенту разницу между tool и resource в MCP." initial_state: ReflectState = { "question": question, "draft": "", @@ -111,23 +95,8 @@ def run_graph(question: str) -> str: "round": 0, "max_rounds": 2, } - 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." - response = await agent.ainvoke( - {"messages": [HumanMessage(content=question)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print("Final answer:\n", response["messages"][-1].content) + final_state = await graph.ainvoke(initial_state) + print("\nFinal Answer:\n", final_state["draft"]) if __name__ == "__main__": asyncio.run(main())