27 lines
756 B
Python
27 lines
756 B
Python
from langgraph.graph import StateGraph, END
|
|
from nodes import draft_answer, reflect, rewrite, ReflectState
|
|
|
|
def build_graph():
|
|
graph = StateGraph(ReflectState)
|
|
|
|
graph.add_node("draft_answer", draft_answer)
|
|
graph.add_node("reflect", reflect)
|
|
graph.add_node("rewrite", rewrite)
|
|
|
|
def condition(state: ReflectState):
|
|
if state["verdict"] == "ok":
|
|
return "ok"
|
|
if state["round"] >= state["max_rounds"]:
|
|
return "maxed"
|
|
return "needs_revision"
|
|
|
|
graph.set_entry_point("draft_answer")
|
|
graph.add_conditional_edges("reflect", condition, {
|
|
"ok": END,
|
|
"needs_revision": "rewrite",
|
|
"maxed": END,
|
|
})
|
|
graph.add_edge("rewrite", "reflect")
|
|
|
|
return graph.compile()
|