my_agent/agent.py

from langgraph.graph import StateGraph, START, END
from .utils.state import State
from .utils.nodes import retrieve, generate_story, generate_brainstroming , generate_final_story, route_after_selection, select_preferred_topics

def build_graph():
    builder = StateGraph(State)
    builder.add_node(retrieve)
    builder.add_node(generate_story)
    builder.add_node(generate_brainstroming)
    builder.add_node(select_preferred_topics)
    builder.add_node(generate_final_story)


    # Normal edges
    builder.add_edge(START, "retrieve")
    builder.add_edge("retrieve", "generate_story")
    builder.add_edge("generate_story", "generate_brainstroming")
    builder.add_edge("generate_brainstroming", "select_preferred_topics")

    # Conditional edge
    builder.add_conditional_edges("select_preferred_topics", route_after_selection,{True:'retrieve',False:'generate_final_story'})
    builder.add_edge("generate_final_story",END)

    return builder.compile()