118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
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())
|