Spaces:
Sleeping
Sleeping
File size: 1,014 Bytes
be3a5c4 ef9fa4b a9f99c3 be3a5c4 a9f99c3 ef9fa4b a9f99c3 ef9fa4b a9f99c3 ef9fa4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from langgraph.graph import StateGraph, START, END
from .utils.state import State
from .utils.nodes import Retriever , ImageCaptioner , StoryGenerator, BrainstromTopicGenerator
from langgraph.checkpoint.memory import MemorySaver
memory = MemorySaver()
class BrainstormingAgent:
def __init__(self):
self.memory = MemorySaver()
def brainstorming_graph(self):
builder = StateGraph(State)
builder.add_node("caption_image",ImageCaptioner().run)
builder.add_node("retrieve",Retriever().run)
builder.add_node("generate_story",StoryGenerator().run)
builder.add_node("generate_brainstroming",BrainstromTopicGenerator().run)
builder.add_edge(START, "caption_image")
builder.add_edge("caption_image", "retrieve")
builder.add_edge("retrieve", "generate_story")
builder.add_edge("generate_story", "generate_brainstroming")
builder.add_edge("generate_brainstroming", END)
return builder.compile(checkpointer=self.memory) |