src/genai/analytics_chatbot/agent.py

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from .utils.state import State
from .utils.nodes import ChatbotNode, FetchDataNode

class ChatbotAgent:
    def __init__(self):
        self.memory = MemorySaver()

    def chatbot_graph(self):
        graph_builder= StateGraph(State)
        graph_builder.add_node("chatbot", ChatbotNode().run)
        graph_builder.add_node("fetch_data", FetchDataNode().run)
        
        graph_builder.add_edge(START, "chatbot")
        graph_builder.add_edge("chatbot", 'fetch_data')
        graph_builder.add_edge("fetch_data", END)


        return graph_builder.compile(checkpointer=self.memory)