diff --git a/main.py b/main.py index 8261a85..5baa944 100644 --- a/main.py +++ b/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([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +backend = CompositeBackend( + default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), + routes={}, +) +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']}" - response = await llm.ainvoke([HumanMessage(content=prompt)]) - state['draft'] = response.content.strip() + prompt = f"Write a concise answer (5–10 sentences) to the following question:\n\n{state['question']}" + response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "draft"}}) + 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)]) - 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 + response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "reflect"}}) + import json + try: + data = json.loads(response["messages"][-1].content) + verdict = data.get("verdict", "needs_revision") + critique = data.get("critique", "") + except Exception: + verdict = "needs_revision" + 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)]) - state['draft'] = response.content.strip() - state['round'] += 1 + response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "rewrite"}}) + new_draft = response["messages"][-1].content + 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.set_entry_point("draft_answer") + graph.add_edge(START, "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) - - app = graph.compile() +async def main(): + question = "Объясни студенту разницу между tool и resource в MCP" initial_state: ReflectState = { "question": question, "draft": "", "critique": "", "verdict": "", "round": 0, - "max_rounds": max_rounds, + "max_rounds": 2, } - final_state = await app.ainvoke(initial_state) - return final_state["draft"] - -# ---------- DeepAgent tool ---------- -@tool -async def answer_question(query: str) -> str: - """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) + graph = build_graph() + result = await graph.ainvoke(initial_state) + print("\n--- Final Answer ---") + print(result["draft"]) + print("\n--- Final Critique ---") + print(result["critique"]) if __name__ == "__main__": asyncio.run(main()) +"""