File size: 2,256 Bytes
3002e1b
 
 
0b2c9fd
3002e1b
3f5fe23
 
 
 
 
 
0b2c9fd
3f5fe23
3002e1b
 
 
 
 
 
3f5fe23
 
 
0b2c9fd
 
3f5fe23
 
 
0b2c9fd
872d043
a6a0614
 
 
b2982ed
 
a6a0614
872d043
93d50e5
872d043
0b2c9fd
 
3002e1b
da7bd22
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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)