langgraph_chatbot_with_tavily_search_and_state.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
10
11# Define the state of our graph
12class State(TypedDict):
13    # Messages are added to the list rather than overwritten
14    messages: Annotated[list, add_messages]
15
16# Initialize the model and tools
17llm = ChatOpenAI(model="gpt-4o-mini")
18tools = [TavilySearchResults(max_results=2)]
19llm_with_tools = llm.bind_tools(tools)
20
21# Define the function that calls the model
22def chatbot(state: State):
23    return {"messages": [llm_with_tools.invoke(state["messages"])]}
24
25# Build the graph
26graph_builder = StateGraph(State)
27
28# Add nodes
29graph_builder.add_node("chatbot", chatbot)
30tool_node = ToolNode(tools=tools)
31graph_builder.add_node("tools", tool_node)
32
33# Add edges
34graph_builder.add_edge(START, "chatbot")
35graph_builder.add_conditional_edges(
36    "chatbot",
37    tools_condition,
38)
39graph_builder.add_edge("tools", "chatbot")
40
41# Compile the graph
42graph = graph_builder.compile()
43
44# Example execution
45def stream_graph_updates(user_input: str):
46    for event in graph.stream({"messages": [("user", user_input)]}):
47        for value in event.values():
48            print("Assistant:", value["messages"][-1].content)
49
50if __name__ == "__main__":
51    # Ensure OPENAI_API_KEY and TAVILY_API_KEY are set in your environment
52    stream_graph_updates("Hi there! Can you help me find the current weather in San Francisco?")