Multi-Agent LangGraph Course – Setting up our multi-agent team

Welcome back to part 5, where we’ll set up our multi-agent team. So buckle up and let’s jump right in. Create a new file named multi_agent.py in your project root:

📂 FINX_LANGGRAPH
    📂 images
    📂 output
    📂 tools
        📄 __init__.py
        📄 image.py
        📄 pdf.py
        📄 weather.py
    📄 .env
    📄 langchain_basics.py
    📄 multi_agent.py    ✨New file
    📄 Pipfile
    📄 Pipfile.lock
    📄 setup_environment.py
    📄 simple_langgraph.py

Open up the multi_agent.py file and start with the imports:

import functools
import operator
from typing import Annotated, Sequence, TypedDict

from colorama import Fore, Style
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph

from setup_environment import set_environment_variables
from tools import generate_image, markdown_to_pdf_file

We have a lot of imports again, many of which will be familiar. We import our own two tools from the tools folder and also the TavilySearchResults from the langchain_community tools. There are some new imports like functools and the AgentExecutor but we’ll cover each one and how they are used as we go along.

Environment variables and constants

Let’s load up our environment variables and create a bunch of constants we’ll need:

set_environment_variables("Multi_Agent_Team")

TRAVEL_AGENT_NAME = "travel_agent"
LANGUAGE_ASSISTANT_NAME = "language_assistant"
VISUALIZER_NAME = "visualizer"
DESIGNER_NAME = "designer"

TEAM_SUPERVISOR_NAME = "team_supervisor"
MEMBERS = [TRAVEL_AGENT_NAME, LANGUAGE_ASSISTANT_NAME, VISUALIZER_NAME]
OPTIONS = ["FINISH"] + MEMBERS

We load our environment variables and set the project name to Multi_Agent_Team. We then define a bunch of constants for the names of our agents and the team supervisor. These are just strings but as we’ll have to type each of these strings multiple times it will be very annoying if we change or mistype one, hence storing these in a single place up top is the way to go.

Note that we have the travel_agent, language_assistant, and visualizer inside a list called members and we have the designer and team_supervisor on the outside. We also imported the END node we used last time. That leaves us with a situation like this:

The list named OPTIONS is going to be the potential options the team_supervisor can choose from each step along the way, so it has all three members in the team + the "FINISH" option to indicate this particular team has finished its work.

Add two more final constants below:

TAVILY_TOOL = TavilySearchResults()
LLM = ChatOpenAI(model="gpt-3.5-turbo-0125")

We have the TAVILY_TOOL which is the Tavily search tool we imported from the langchain_community tools and the LLM which is gpt-3.5-turbo-0125 here but feel free to use GPT-4-turbo instead if you want.

Agent creator function

We’re going to be creating a lot of agents here, so let’s create a function to handle the repetitive work of creating an agent:

def create_agent(llm: BaseChatModel, tools: list, system_prompt: str):
    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    agent = create_openai_tools_agent(llm, tools, prompt_template)
    agent_executor = AgentExecutor(agent=agent, tools=tools)  # type: ignore
    return agent_executor

We define a function named create_agent which takes an llm of the type BaseChatModel. This is just a type hint but it was part of our imports for clarity. BaseChatModel is the base class for all chat models in LangChain, including the ChatOpenAI variation we use here. You can pass any LLM you want and have different nodes of the same graph run on completely different LLMs. The other arguments are a list of tools and a system_prompt string.

We then declare a prompt_template using the ChatPromptTemplate.from_messages method that we used all the way back in part 1, but this time we use multiple messages. We have a "system" message that is the system prompt string passed into the function and then we have two placeholders for the messages and agent_scratchpad variables that we have seen before. The MessagesPlaceholder, as the name suggests, is just a placeholder for both of these so we can insert them later using the names we have defined under variable_name.

We then use the create_openai_tools_agent just like we did in part 3, but this time we go one step further and create an AgentExecutor in the step below. This AgentExecutor comes with LangChain and will basically combine the agent and the executor nodes we had in the previous part into a single node, handling the function call logic we did in the previous part for us! It takes an agent and a list of tools for that agent to use as arguments.

The # type: ignore comment is in case you use a type checker as it will complain here, and this series is not about type checking so we won’t go too deep into it as it’s no big deal here. We then return the agent_executor we created.

Agent state object

Now let’s declare the state object that we will be passing around in this particular graph:

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str

This time we need two entries. The first is the messages which is a sequence of BaseMessage objects which again are just messages like ("human", "Hello, how are you doing?"), or ("ai", "I'm doing well, thanks!"),. We define it as a Sequence, so like a list or a tuple of these messages, and the operator.add again indicates that we will add to this sequence of messages with each step. Annotated is just used as it allows us to add the annotation of operator.add.

The second entry is the next which is a string that will be the name of the next agent to call. This is the agent that the team_supervisor will decide to call next based on the state object it receives and then we can use this field to see which agent to route to next. This field can just be overwritten as we don’t need the history, so a single string without any fancy annotations will do fine here.

Agent node function

Now let’s define a function that represents one of these agent nodes in general:

def agent_node(state, agent, name):
    result = agent.invoke(state)
    return {"messages": [HumanMessage(content=result["output"], name=name)]}

The function takes the state object, an agent, and the string name for the agent (the ones we defined up top as constants). Then we simply need to invoke the agent with the state and then keeping with the promise we made above in the AgentState object we defined the node needs to return a messages object with a message in it. We will simply use a HumanMessage, as it doesn’t really matter who the message comes from, and get the result from result["output"] which is the output of the agent’s call.

Team supervisor’s next member choice

Next, we’re going to need a way to have the team_supervisor choose which agent to invoke next. The easiest way to do this reliably is to pretend this is a function that the agent supervisor has to call for us. The only possible input arguments are the names of our agents and we tell the team_supervisor that it must call nonexistent_function(agent_name) to invoke the agent.

This is a bit of a hack, but it makes it very easy for us to extract the agent_name consistently and easily to see which agent node needs to run next. We will also include one extra option of “FINISH” so the team_supervisor can tell us when it’s done and needs to break out of the team. Doing this will also let us use the JsonOutputFunctionsParser later on in our code, as the function call will be sent in a correct JSON format, making the parsing of the output easier.

For this function that doesn’t actually exist, we’re going to define an old-school vanilla OpenAI function description that describes how the function works to the LLM team supervisor. Add the following variable:

router_function_def = {
    "name": "route",
    "description": "Select the next role.",
    "parameters": {
        "title": "routeSchema",
        "type": "object",
        "properties": {
            "next": {
                "title": "next",
                "anyOf": [
                    {"enum": OPTIONS},
                ],
            }
        },
        "required": ["next"],
    },
}

This is actually JSON Schema vocabulary, but is quite readable. We define the name of the function as route and give it a description of what the function does. We then define the parameters that the function takes, giving the parameter object a title of routeSchema and defining that it is an object. Then we define the properties of this object, which is just a single property named next. This property has a title of next and the options available are anyOf the enumerate (list) of OPTIONS we defined up top. We then define that the next property is required.

This JSON Schema style is what the OpenAI API normally uses for function/tool calls, but LangChain has done this under the hood for the functions we have used so far. Again, this function will not actually exist, but that doesn’t stop us from feeding it to the LLM and extracting the next property from the arguments the LLM provides for us.

Team supervisor system prompt

Now let’s create a secondary file to store our prompt system setup messages as we’re going to be using quite a lot of them here. Create a new file named multi_agent_prompts.py in your project root:

📂 FINX_LANGGRAPH
    📂 images
    📂 output
    📂 tools
        📄 __init__.py
        📄 image.py
        📄 pdf.py
        📄 weather.py
    📄 .env
    📄 langchain_basics.py
    📄 multi_agent.py
    📄 multi_agent_prompts.py    ✨New file
    📄 Pipfile
    📄 Pipfile.lock
    📄 setup_environment.py
    📄 simple_langgraph.py

We’ll use this file to store the prompt string variables for the system messages our agents will use. If you’re watching the video tutorial version of this please be advised that there is a written blog version of this tutorial where you can copy these prompts so you don’t have to type them all over again, as we have a lot more of them coming. Let’s start with the team supervisor. Inside the multi_agent_prompts.py file add:

TEAM_SUPERVISOR_SYSTEM_PROMPT = """
You are a supervisor tasked with managing a conversation between the following workers: {members}. Given the following user request, respond with the worker to act next. Each worker will perform a task and respond with their results and status. The end goal is to provide a good travel itinerary for the user, with things to see and do, practical tips on how to deal with language difficulties, and a nice visualization that goes with the travel plan (in the form of an image path, the visualizer will save the image for you and you only need the path).

Make sure you call on each team member ({members}) at least once. Do not call the visualizer again if you've already received an image file path. Do not call any team member a second time unless they didn't provide enough details or a valid response and you need them to redo their work. When finished, respond with FINISH, but before you do, make sure you have a travel itinerary, language tips for the location, and an image file-path. If you don't have all of these, call the appropriate team member to get the missing information.
"""

So we have some basic instructions for the team supervisor on how to manage the team here. We have the placeholder {members} in there twice which will be replaced with the actual list of members. We tell it we want a travel itinerary with things to do and sightseeing, language tips, and a visualization for the itinerary. The prompt here is far from perfect and you can tweak it further if you like.

Save the multi_agent_prompts.py file and let’s get back to the multi_agent.py file. First of all, add an extra import up top with the other imports:

#... all the other imports ...

from multi_agent_prompts import TEAM_SUPERVISOR_SYSTEM_PROMPT

Note that we could just use from multi_agent_prompts import * as the * will simply import everything from the file, even the variables we add later, but this is a bad practice as it makes it hard to see where the variables come from and leads to namespace pollution. It’s better to explicitly define and keep track of what you’re importing or sooner or later you’re going to have multiple variables with the same name and you won’t know where they come from.

Team supervisor prompt template

Now scroll all the way back down past the router_function_def and add the following code to define our team supervisor’s prompt template manually as it will be different from all the other agents:

team_supervisor_prompt_template =