AI Agents with LangGraph

Transcript

Tools Demo

Open the tools.ipynb notebook in the Starter folder. To get your feet wet working with tools, you’ll use Tavily search from the LangChain community. Begin by installing the necessary libraries:

!pip install langchain-community tavily-python

Go to the website at tavily.com, sign up, and get an API key. Once you have the key, store it in your .env file with the name TAVILY_API_KEY. You should also have your OpenAI API key in there as well since you’ll need it later in the lesson:

OPENAI_API_KEY=<your-api-key>
TAVILY_API_KEY=<your-api-key>

Then load the API keys into your environment:

from dotenv import load_dotenv
load_dotenv()

Import the tool from the community library and instantiate it:

from langchain_community.tools.tavily_search import TavilySearchResults

tool = TavilySearchResults()

Now, perform a search just using the tool:

tool.invoke({"query": "What's in the AI news?"})

The results give you up-to-date information. Since LLMs use pre-trained models from the past, it’s very useful to get current information like this to integrate into your AI Agent.

Next, import the ChatOpenAI class to use as your LLM. This model can handle more than one tool, so even though you only have one tool for now, put it in a list and then bind it to the LLM:

import os
from langchain_openai import ChatOpenAI

tools = [tool]

llm = ChatOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
llm_with_tools = llm.bind_tools(tools)

Calling bind_tools is all you need to do to let the model know about your tool. It won’t invoke the tool itself, but when the model encounters a situation in which it recognizes the tool would be helpful, it’ll reply with a message saying which tool to use.

You’ll save the list of human, AI, and tool messages in a list within a state object. Define that class now:

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
import operator

class State(TypedDict):
  messages: Annotated[Sequence[BaseMessage], operator.add]

As you probably recall, a BaseMessage could be a HumanMessage, AIMessage, or ToolMessage. The operator.add tells LangGraph to append new chat messages to the list of messages as the state travels through the graph.

Define a function to call the LLM to respond to human or tool input:

def call_llm(state):
  messages = state["messages"]
  response = llm_with_tools.invoke(messages)
  return {"messages": [response]}

The LLM takes the message list as input and then adds its own AIMessage response to the list. This response could be a normal chat response or a request to use a tool.

Next, prepare your graph. Since you’re passing around state, use a StateGraph with the custom State that you defined earlier:

from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode, tools_condition

graph = StateGraph(State)

graph.add_node("chatbot", call_llm)
graph.add_node("tools", ToolNode(tools))

graph.add_edge(START, "chatbot")
graph.add_conditional_edges(
  "chatbot", tools_condition
)
graph.add_edge("tools", "chatbot")

app = graph.compile()

You added two nodes, one for the LLM, which you’re calling chatbot, and one for the tools. ToolNode is a special node that recognizes if the previous AIMessage calls for a tool. If it does, ToolNode will invoke that tool. You have a conditional edge from the chatbot. tools_condition is a special routing function that will route to the tools node if the AIMessage coming from the chatbot is a tool call or will route to END if it’s not. For this to work, the tool node must be named "tools". Finally, you have a normal edge pointing back from the tools node to the chatbot. This allows LangGraph to keep looking until the LLM has enough information to respond.

Check out the structure:

from IPython.display import Image, display

display(Image(app.get_graph().draw_mermaid_png()))

You can see the loop between the chatbot and the tools.

Finish off the code by invoking your graph with a human message:

inputs = {"messages": [HumanMessage(content="What's the latest AI news?")]}
app.invoke(inputs)

Now check the list of messages. Here’s the HumanMessage. And here you have the AIMessage. The content is empty, but down here, you have tool_calls for tavily_search_results_json. That triggers your ToolNode to invoke the search tool and return a ToolMessage. When the ToolMessage comes back, it’s routed back to the chatbot. Now the chatbot has enough information to respond: “Here are some of the latest AI news articles”.

What if you change the input message to “How many letters are in asjdfkasjdlfajsd?”

The LLM tried using the web search tool, but it turns out that it isn’t a great tool for counting letters. You can create your own tool to do that.

Comment out TavilySearchResults and the tool.invoke and add this function:

from langchain_core.tools import tool

@tool
def count_characters(text: str) -> int:
  """Counts the number of characters in the text"""
  return len(text)

tool = count_characters

Rerun all the cells. Now, the LLM calls your new tool! How did it know? It’s based on the context you gave it in the docstring and other tool information.