src/genai/brainstroming_agent/agent.py

from langgraph.graph import StateGraph, START, END
from .utils.state import State
from .utils.nodes import retrieve, generate_story, generate_brainstroming,caption_image
from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()


def brainstroming_graph():
    builder = StateGraph(State)
    builder.add_node(caption_image)
    builder.add_node(retrieve)
    builder.add_node(generate_story)
    builder.add_node(generate_brainstroming)


    # Normal edges
    builder.add_edge(START, "caption_image")
    builder.add_edge("caption_image", "retrieve")

    # builder.add_edge(START, "retrieve")
    builder.add_edge("retrieve", "generate_story")

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