From 9500364d27dd8933ce367da46f08bddcfb3f97eb 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: Thu, 4 Jun 2026 16:26:36 +0000 Subject: [PATCH] add: main.py --- main.py | 207 ++++++++++++++++++++++++-------------------------------- 1 file changed, 89 insertions(+), 118 deletions(-) diff --git a/main.py b/main.py index 6fbaea0..922a39b 100644 --- a/main.py +++ b/main.py @@ -1,22 +1,16 @@ -""" -# main.py -# Implementation of the LangGraph reflective agent using deepagents -# Author: Auto-generated for the assignment -# Requires: deepagents, langchain-openai, langgraph, langchain -""" import os import asyncio from typing import TypedDict, Annotated from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.messages import HumanMessage from langchain.tools import tool from deepagents import create_deep_agent -from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from langgraph.graph import StateGraph, START, END from langgraph.graph.message import add_messages -# ---------- LLM configuration (OpenRouter) ---------- +# ---------- LLM ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -24,130 +18,107 @@ llm = ChatOpenAI( 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 implementations ---------- -async def draft_answer(state: ReflectState) -> ReflectState: - """Generate an initial draft answer (5–10 sentences).""" - prompt = ( - "You are an expert tutor. Answer the following question in 5–10 concise sentences. " - "Avoid filler and keep it clear. -" - f"Question: {state['question']}" - ) - response = await llm.ainvoke([SystemMessage(content="You are a helpful tutor."), HumanMessage(content=prompt)]) - state["draft"] = response.content.strip() - state["round"] = 1 - return state - -async def reflect(state: ReflectState) -> ReflectState: - """Critique the draft and decide if revision is needed.""" - prompt = ( - "You are a critical reviewer. Evaluate the following answer for completeness, specificity, and lack of filler. " - "Respond with a verdict of either "ok" or "needs_revision", followed by 2–3 bullet points of constructive feedback. -" - f"Answer draft:\n{state['draft']}" - ) - response = await llm.ainvoke([SystemMessage(content="You are a critical reviewer."), HumanMessage(content=prompt)]) - text = response.content.strip() - # Parse verdict and critique - if "needs_revision" in text.lower(): - verdict = "needs_revision" - else: - verdict = "ok" - # Extract critique lines after the verdict - lines = text.splitlines() - critique_lines = [line for line in lines if line.strip() and line.strip().lower() not in {"ok", "needs_revision"}] - critique = "\n".join(critique_lines).strip() - state["critique"] = critique - state["verdict"] = verdict - return state - -async def rewrite(state: ReflectState) -> ReflectState: - """Rewrite the draft incorporating the critique.""" - prompt = ( - "You are revising an answer based on the following critique. Produce a new version that addresses the points and remains 5–10 sentences. -" - f"Original draft:\n{state['draft']}\n\nCritique:\n{state['critique']}" - ) - response = await llm.ainvoke([SystemMessage(content="You are a revising tutor."), HumanMessage(content=prompt)]) - state["draft"] = response.content.strip() - state["round"] += 1 - return state - -# ---------- Graph construction ---------- -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) - - # Entry point - graph.set_entry_point("draft_answer") - - # Transitions - graph.add_conditional_edges( - "draft_answer", - lambda _: "reflect", - ) - graph.add_conditional_edges( - "reflect", - lambda state: "END" if state["verdict"] == "ok" else "rewrite", - ) - graph.add_conditional_edges( - "rewrite", - lambda state: "END" if state["round"] > state["max_rounds"] else "reflect", - ) - - return graph - -# ---------- Tool that runs the graph ---------- -@tool -async def answer_question(query: str) -> str: - """Run the reflective LangGraph to answer a question.""" - graph = build_graph() - # Initialize state - state: ReflectState = { - "question": query, - "draft": "", - "critique": "", - "verdict": "", - "round": 0, - "max_rounds": 2, - } - # Run graph - final_state = await graph.ainvoke(state) - return final_state["draft"] - -# ---------- DeepAgent setup ---------- +# ---------- Backend for deepagents ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) +# ---------- State definition ---------- +class ReflectState(TypedDict): + question: str + draft: str + critique: str + verdict: str # ok | needs_revision + round: int + max_rounds: int + +# ---------- LangGraph 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 = ( + "You are a critic. Evaluate the draft answer for completeness, specificity, and lack of filler.\n" + "Return a verdict ('ok' or 'needs_revision') and 2–3 bullet points of critique.\n" + f"Draft: {state['draft']}" + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + # Simple parsing: first line verdict, rest critique + lines = response.content.strip().splitlines() + verdict = lines[0].strip().lower() + critique = "\n".join(lines[1:]).strip() + state["verdict"] = verdict + state["critique"] = critique + return state + +async def rewrite(state: ReflectState) -> ReflectState: + prompt = ( + "Rewrite the draft answer incorporating the following critique. Keep the answer concise (5–10 sentences).\n" + f"Critique: {state['critique']}\n" + f"Original draft: {state['draft']}" + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state["draft"] = response.content.strip() + state["round"] += 1 + return state + +# ---------- Graph construction ---------- +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") + +# Transitions +graph.add_edge("draft_answer", "reflect") +# From reflect: if ok -> END, else if needs_revision and round < max_rounds -> rewrite +graph.add_conditional_edges( + "reflect", + lambda state: ( + "END" if state["verdict"] == "ok" else ( + "rewrite" if state["round"] < state["max_rounds"] else "END" + ) + ), +) +graph.add_edge("rewrite", "reflect") + +app = graph.compile() + +# ---------- DeepAgent wrapper ---------- +@tool +def run_graph(question: str, max_rounds: int = 2) -> str: + """Run the reflection graph for a given question.""" + initial_state: ReflectState = { + "question": question, + "draft": "", + "critique": "", + "verdict": "", + "round": 0, + "max_rounds": max_rounds, + } + result = app.invoke(initial_state) + return result["draft"] + agent = create_deep_agent( model=llm, - tools=[answer_question], + tools=[run_graph], backend=backend, - system_prompt="You are a helpful educational agent. Use the provided tools to answer questions.", + system_prompt="You are an assistant that can generate and improve answers using reflection. Use the provided tool to get a refined answer.", ) -# ---------- CLI entry point ---------- async def main(): question = "Объясни студенту разницу между tool и resource в MCP" - result = await agent.ainvoke( + response = await agent.ainvoke( {"messages": [HumanMessage(content=question)]}, {"configurable": {"thread_id": "session-1"}}, ) - # The tool returns the final answer as the last message content - print("\nFinal answer:\n", result["messages"][-1].content) + print("Final answer:\n", response["messages"][-1].content) if __name__ == "__main__": asyncio.run(main())