From 1043b6fc8b79d8153f9f192b01faef988fa93bd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Sat, 27 Jun 2026 13:49:58 +0000 Subject: [PATCH] add: main.py --- main.py | 117 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..9381a3c --- /dev/null +++ b/main.py @@ -0,0 +1,117 @@ +import os +import json +import asyncio +from typing import TypedDict + +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend +from langgraph.graph import StateGraph, START, END + +# ---------- LLM ---------- +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, +) + +# ---------- Backend & Agent ---------- +backend = FilesystemBackend() +agent = create_deep_agent( + model=llm, + tools=[], + backend=backend, + system_prompt="You are a helpful assistant.", +) + +# ---------- 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 brief answer (5-10 sentences) to the following question: {state['question']}" + result = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "draft"}}, + {}, + ) + state["draft"] = result["messages"][-1].content + return state + +async def reflect(state: ReflectState) -> ReflectState: + prompt = ( + f"Critique the following answer. Provide verdict ok or needs_revision and 2-3 points of critique in JSON format with keys verdict and critique.\nAnswer: {state['draft']}" + ) + result = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "reflect"}}, + {}, + ) + content = result["messages"][-1].content + try: + data = json.loads(content) + state["verdict"] = data.get("verdict", "").lower() + state["critique"] = data.get("critique", "") + except json.JSONDecodeError: + state["verdict"] = "needs_revision" + state["critique"] = content + return state + +async def rewrite(state: ReflectState) -> ReflectState: + prompt = ( + f"Rewrite the answer to improve it based on the critique: {state['critique']}\nPrevious draft: {state['draft']}\nProvide the improved answer." + ) + result = await agent.ainvoke( + {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "rewrite"}}, + {}, + ) + state["draft"] = result["messages"][-1].content + state["round"] += 1 + return state + +# ---------- Conditional Edge ---------- +def reflect_cond(state: ReflectState): + if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]: + return "rewrite" + return "end" + +# ---------- Graph ---------- + +graph = StateGraph(ReflectState) +graph.add_node("draft_answer", draft_answer) +graph.add_node("reflect", reflect) +graph.add_node("rewrite", rewrite) + +graph.add_edge(START, "draft_answer") +graph.add_edge("draft_answer", "reflect") +graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", "end": END}) +graph.add_edge("rewrite", "reflect") + +graph.set_entry_point("draft_answer") +graph.set_finish_point(END) + +executor = graph.compile() + +# ---------- CLI ---------- +async def main(): + question = input("Enter a question: ") + initial_state: ReflectState = { + "question": question, + "draft": "", + "critique": "", + "verdict": "", + "round": 1, + "max_rounds": 2, + } + final_state = await executor(initial_state) + print("\nFinal answer:\n") + print(final_state["draft"]) + +if __name__ == "__main__": + asyncio.run(main())