langgraph_chatbot_with_tavily_search_and_memory_checkpoint.py

1import os
2from typing import Annotated
3from typing_extensions import TypedDict
4
5from langchain_openai import ChatOpenAI
6from langchain_community.tools.tavily_search import TavilySearchResults
7from langgraph.graph import StateGraph, START, END
8from langgraph.graph.message import add_messages
9from langgraph.prebuilt import ToolNode, tools_condition
10from langgraph.checkpoint.memory import MemorySaver
11
12# 1. Define the State
13class State(TypedDict):
14    # add_messages ensures new messages are appended to the history rather than overwriting it
15    messages: Annotated[list, add_messages]
16
17# 2. Initialize the Graph
18graph_builder = StateGraph(State)
19
20# 3. Define the Tools and LLM
21tool = TavilySearchResults(max_results=2)
22tools = [tool]
23llm = ChatOpenAI(model="gpt-4o")
24llm_with_tools = llm.bind_tools(tools)
25
26# 4. Define the Node functions
27def chatbot(state: State):
28    return {"messages": [llm_with_tools.invoke(state["messages"])]}
29
30# 5. Build the Graph
31graph_builder.add_node("chatbot", chatbot)
32
33tool_node = ToolNode(tools=[tool])
34graph_builder.add_node("tools", tool_node)
35
36# The tools_condition helper defines the path from chatbot to tools or END
37graph_builder.add_conditional_edges(
38    "chatbot",
39    tools_condition,
40)
41graph_builder.add_edge("tools", "chatbot")
42graph_builder.add_edge(START, "chatbot")
43
44# 6. Add Memory and Compile
45memory = MemorySaver()
46graph = graph_builder.compile(checkpointer=memory)
47
48# 7. Run the Graph
49config = {"configurable": {"thread_id": "1"}}
50
51def stream_graph_updates(user_input: str):
52    events = graph.stream(
53        {"messages": [("user", user_input)]}, config, stream_mode="values"
54    )
55    for event in events:
56        event["messages"][-1].pretty_print()
57
58if __name__ == "__main__":
59    # Example usage:
60    stream_graph_updates("Hi! I'm Bob.")
61    stream_graph_updates("What's my name?")