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,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)
    # 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")

    # Conditional edge
    builder.add_edge("generate_brainstroming", END)
    # builder.add_edge("generate_final_story",END)

    return builder.compile(checkpointer=memory)