updated the nodes in agent

my_agent/agent.py

@@ -1,39 +1,10 @@
 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,caption_image
+from .utils.nodes import retrieve, generate_story, generate_brainstroming , route_after_selection, select_preferred_topics,caption_image
 from langgraph.checkpoint.memory import MemorySaver
 
 memory = MemorySaver()
 
-def build_graph_old():
-    builder = StateGraph(State)
-    # builder.add_node(caption_image)
-    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, "caption_image")
-
-    builder.add_edge(START, "retrieve")
-    builder.add_edge("retrieve", "generate_story")
-
-    # builder.add_edge("caption_image", "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(checkpointer=memory)
-
-
-
 
 def build_graph():
     builder = StateGraph(State)
@@ -42,7 +13,6 @@ def build_graph():
     builder.add_node(generate_story)
     builder.add_node(generate_brainstroming)
     builder.add_node(select_preferred_topics)
-    builder.add_node(generate_final_story)
 
 
     # Normal edges