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 FetchDataNode , FetchLastMessage , RetrievePossibleEndpoints , FetchParametersNode , RetrieveExactEndpoint, BackupRetrievalNode, BackupRoutingNode, QueryCheckNode

fetch_last_message_node = FetchLastMessage()
retrieve_api_endpoints_node = RetrievePossibleEndpoints()
retrieve_exact_endpoint_node = RetrieveExactEndpoint()
fetch_parameters_node = FetchParametersNode()
fetch_data_node = FetchDataNode()
backup_response_node = BackupRetrievalNode()
query_check_node = QueryCheckNode()

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

    def chatbot_graph(self):
        graph_builder= StateGraph(State)
        graph_builder.add_node("fetch_last_message", fetch_last_message_node.run)
        graph_builder.add_node("retrieve_api_endpoints", retrieve_api_endpoints_node.run)
        graph_builder.add_node("retrieve_exact_endpoint", retrieve_exact_endpoint_node.run)
        graph_builder.add_node("check_query_type", query_check_node.run)

        graph_builder.add_node("fetch_parameters", fetch_parameters_node.run)
        graph_builder.add_node("fetch_data", fetch_data_node.run)
        graph_builder.add_node("backup_response", backup_response_node.run)

        
        graph_builder.add_edge(START, "fetch_last_message")
        graph_builder.add_edge("fetch_last_message", 'retrieve_api_endpoints')
        graph_builder.add_edge("retrieve_api_endpoints", 'retrieve_exact_endpoint')
        graph_builder.add_edge("retrieve_exact_endpoint", 'fetch_parameters')
        # graph_builder.add_edge("check_query_type", 'fetch_parameters')
        graph_builder.add_edge("fetch_parameters", 'fetch_data')
        graph_builder.add_edge("fetch_data", END)
        graph_builder.add_edge("backup_response", END)

        graph_builder.add_conditional_edges("fetch_parameters", BackupRoutingNode().run,{'execute_backup':'backup_response', 'go_on':"fetch_data"})
        graph_builder.add_conditional_edges("fetch_data", BackupRoutingNode().run,{'execute_backup':'backup_response', 'go_on':END})

        return graph_builder.compile(checkpointer=self.memory)