tool_calling_graph.py

from typing import Annotated
from typing_extensions import TypedDict
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langchain_tavily import TavilySearch
from langgraph.prebuilt import ToolNode, tools_condition

# Load environment variables from a .env file if present
load_dotenv()


# Create a Graph
class State(TypedDict):
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)

llm = init_chat_model("deepseek:deepseek-chat")

tool = TavilySearch(max_results=2)
tools = [tool]

# Tell the LLM which tools it can call
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State) -> State:
    messages = llm_with_tools.invoke(state["messages"])
    return {"messages": [messages]}


graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=[tool])
graph_builder.add_node("tools", tool_node)

# The `tools_condition` function returns "tools" if the chatbot asks to use a tool, and "END" if
# it is fine directly responding. This conditional routing defines the main agent loop.
graph_builder.add_conditional_edges("chatbot", tools_condition)

# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")

# Add an entry point
graph_builder.add_edge(START, "chatbot")

# Compile the graph
graph = graph_builder.compile()