add: main.py

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())