49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
# imports from langgraph
|
|
from langgraph.graph import START, END
|
|
|
|
# imports from local files
|
|
from src.graph_builder import graph_builder
|
|
from src.nodes import chatbot, BasicToolNode
|
|
from src.utils import route_tools
|
|
from src.tools import tool
|
|
|
|
|
|
def main():
|
|
# The first argument is the unique node name
|
|
# The second argument is the function or object that will be called whenever
|
|
# the node is used.
|
|
graph_builder.add_node("chatbot", chatbot) # Create the chatbot node
|
|
graph_builder.add_edge("chatbot", END) # Add the edge from chatbot node to END
|
|
tool_node = BasicToolNode(tools=[tool]) # Intermediate variable creation to create a tool node with the tools
|
|
graph_builder.add_node("tools", tool_node) # Create a tool node with the tools
|
|
graph_builder.add_conditional_edges( # Create the conditionnal edges from chatbot to tools
|
|
"chatbot",
|
|
route_tools,
|
|
{"tools": "tools", END: END},
|
|
)
|
|
graph_builder.add_edge("tools", "chatbot") # Add the edge from tools node back to chatbot
|
|
graph_builder.add_edge(START, "chatbot") # Add the edge from START to chatbot node
|
|
graph = graph_builder.compile()
|
|
|
|
def stream_graph_updates(user_input: str):
|
|
for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
|
|
for value in event.values():
|
|
print("Assistant:", value["messages"][-1].content)
|
|
|
|
|
|
while True:
|
|
try:
|
|
user_input = input("User: ")
|
|
if user_input.lower() in ["quit", "exit", "q"]:
|
|
print("Goodbye!")
|
|
break
|
|
stream_graph_updates(user_input)
|
|
except:
|
|
# fallback if input() is not available
|
|
user_input = "What do you know about LangGraph?"
|
|
print("User: " + user_input)
|
|
stream_graph_updates(user_input)
|
|
break
|
|
|
|
if __name__ == "__main__":
|
|
main() |