| |
| |
|
|
| from langgraph.graph import StateGraph, START, END |
| from langgraph.checkpoint.memory import MemorySaver |
| from nodes import ( |
| State, |
| speaker_chatbot_node, |
| tool_node |
| ) |
| from langgraph.prebuilt import tools_condition |
| from form_management.form_management import Form |
| import sys |
| from IPython.display import Image, display |
| import logging |
|
|
| class GraphManager: |
| def __init__(self): |
| self.graph = None |
| self.init_graph() |
|
|
| |
| |
| |
| def init_graph(self): |
| graph_memory = MemorySaver() |
| graph_builder = StateGraph(State) |
|
|
| graph_builder.add_node("speaker_chatbot", speaker_chatbot_node) |
| graph_builder.add_node("tools", tool_node) |
|
|
| graph_builder.add_conditional_edges( |
| "speaker_chatbot", |
| tools_condition, |
| ) |
| |
| graph_builder.add_edge("tools", "speaker_chatbot") |
|
|
| graph_builder.add_edge(START, "speaker_chatbot") |
| graph_builder.add_edge("speaker_chatbot", END) |
|
|
| self.graph = graph_builder.compile(checkpointer=graph_memory) |
|
|
| def save_graph_image(self): |
| """ |
| Save the graph as a PNG image. |
| """ |
| image = Image(self.graph.get_graph().draw_mermaid_png()) |
| with open("graph.png", "wb") as f: |
| f.write(image.data) |
|
|
| def process_user_input(self, user_input, session_id): |
| """ |
| Process the user input and update the graph state. |
| """ |
| config = {"configurable": {"thread_id": session_id}} |
|
|
| input_dict = { |
| "messages": [("user", user_input)], |
| "session_id": session_id |
| } |
|
|
| for event in self.graph.stream( |
| input_dict, |
| config, |
| stream_mode="values" |
| ): |
| event["messages"][-1].pretty_print() |
|
|
| snapshot = self.graph.get_state(config) |
| return snapshot[0]["messages"][-1].content |
| |
| def user_input_handler(user_input, session_id): |
| logging.info(user_input) |
| response = graph_manager.process_user_input(user_input, session_id) |
| return response |
|
|
| graph_manager = GraphManager() |
| graph_manager.save_graph_image() |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
