File size: 960 Bytes
9cccf74 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
from langgraph.graph import StateGraph
from builder.state import AgentState
from builder.nodes import sql_generator, sql_executor, chart_generator, narrator, decider, general_chat
def build_graph():
graph = StateGraph(AgentState)
graph.add_node("decider", decider)
graph.add_node("sql_generator", sql_generator)
graph.add_node("sql_executor", sql_executor)
graph.add_node("chart_generator", chart_generator)
graph.add_node("narrator", narrator)
graph.add_node("general_chat", general_chat)
graph.set_entry_point("decider")
graph.add_conditional_edges(
"decider",
lambda state: state["route"],
{
"sql": "sql_generator",
"chat": "general_chat",
},
)
graph.add_edge("sql_generator", "sql_executor")
graph.add_edge("sql_executor", "chart_generator")
graph.add_edge("sql_executor", "narrator")
return graph.compile()
|