from langgraph.graph import StateGraph,START,END
from .state import TravelState
from .agents import transport_agent_node,hotel_agent_node,attraction_agent_node,coordinator_node,integrator_agent_node

def build_graph():
    builder=StateGraph(TravelState)
    #添加节点
    builder.add_node("coordinator",coordinator_node)
    builder.add_node("transport",transport_agent_node)
    builder.add_node("hotel",hotel_agent_node)
    builder.add_node("attraction",attraction_agent_node)
    builder.add_node("integrator",integrator_agent_node)
    #添加边(并行)
    builder.add_edge(START,"coordinator")
    builder.add_edge("coordinator","transport")
    builder.add_edge("coordinator","hotel")
    builder.add_edge("coordinator","attraction")

    builder.add_edge("transport","integrator")
    builder.add_edge("hotel","integrator")
    builder.add_edge("attraction","integrator")
    builder.add_edge("integrator",END)

    #编译图
    graph= builder.compile()

    return graph

async def main():
    graph=build_graph()
    result=await graph.ainvoke({"user_request":"成都三日游，两人，喜欢美食和人文"})
    print(result["final_plan"])

    
