langgraph/examples/advanced_agents/multi-agent/multi-agent-collaboration.ipynb

Basic Multi-agent Collaboration

A single agent can usually operate effectively using a handful of tools and a scoped domain, but even using powerful models like gpt-4, it can be less effective at using many tools. One way to improve your system's performance is through multi-agent collaboration. Each agent can specialize in a task or domain, and LangGraph can effectively orchestrate them to accomplish a larger goal.

This notebook is inspired by the paper AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation, by Wu, et. al.

This and other "advanced agent" notebooks are designed to show how you can implement certain design patterns in LangGraph. If the pattern suits your needs, we recommend combining it with some of the other fundamental patterns described elsewhere in the docs for best performance.

In [ ]:

%pip install -U langchain langchain_openai langsmith pandas

In [ ]:

import getpass
import os


def _set_if_undefined(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass(f"Please provide your {var}")


_set_if_undefined("OPENAI_API_KEY")
_set_if_undefined("LANGCHAIN_API_KEY")
_set_if_undefined("TAVILY_API_KEY")

# Optional, add tracing in LangSmith
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Multi-agent Collaboration"

Create tool

Below is an example of 2 agents collaborating to accomplish a single task.

In [ ]:

from typing import List, Tuple, Union

import matplotlib.pyplot as plt
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool

tavily_tool = TavilySearchResults(max_results=5)


@tool
def create_plot(
    data: Union[List[float], List[int]],
    labels: Union[List[str], None] = None,
    title: str = "Plot",
    xlabel: str = "X",
    ylabel: str = "Y",
    color: Union[str, List[str]] = "blue",
    plot_type: str = "bar",
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Generates a bar or line plot from the provided data and returns the figure and axis objects.

    :param data: A list of numerical values for the bar heights or line points.
    :param labels: A list of strings for the bar or point labels. Default is None.
    :param title: Title of the plot. Default is 'Plot'.
    :param xlabel: Label for the X-axis. Default is 'X'.
    :param ylabel: Label for the Y-axis. Default is 'Y'.
    :param color: Color of the bars or line. Can be a single color or a list of colors. Default is 'blue'.
    :param figsize: Size of the figure as a tuple (width, height). Default is (10, 6).
    :param plot_type: Type of plot ('bar' or 'line'). Default is 'bar'.
    :return: Tuple containing the figure and axes objects.
    """
    if plot_type not in ["bar", "line"]:
        raise ValueError("Invalid plot_type. Expected 'bar' or 'line'.")

    fig, ax = plt.subplots(figsize=(10, 6))
    x_positions = range(len(data))

    if labels and len(labels) == len(data):
        plt.xticks(x_positions, labels)

    if plot_type == "bar":
        ax.bar(x_positions, data, color=color)
    elif plot_type == "line":
        ax.plot(x_positions, data, color=color, marker="o")  # 'o' for circular markers

    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    return fig, ax

Create the graph

We will create the individual agents below and place them within the graph.

In [ ]:

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

from langchain.agents import create_openai_functions_agent
from langchain.tools.render import format_tool_to_openai_function
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    ChatMessage,
    FunctionMessage,
    HumanMessage,
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from typing_extensions import TypedDict

from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation


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


workflow = StateGraph(AgentState)

# This a helper class we have that is useful for running tools
# It takes in an agent action and calls that tool and returns the result
tools = [tavily_tool, create_plot]
tool_executor = ToolExecutor(tools)

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful AI assistant, collaborating with other assistants."
            " Use the provided tools to progress towards answering the question."
            " If you are unable to fully answer, that's OK, another assistant with different tools "
            " will help where you left off. Execute what you can to make progress."
            " If you have the final answer, prefix your response with FINAL ANSWER."
            " You have access to the following tools: {tool_names}.",
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)


def add_agent_node(name: str, llm, tools, prompt):
    functions = [format_tool_to_openai_function(t) for t in tools]

    def _update_state(ai_message) -> dict:
        if isinstance(ai_message, FunctionMessage):
            result = ai_message
        else:
            message = ai_message.dict(exclude={"type"})
            message["name"] = name
            result = HumanMessage(**message)
        return {
            "messages": [result],
            "sender": name,
        }

    chain = (
        (lambda x: {**x, "intermediate_steps": []})
        | prompt.partial(tool_names=", ".join([tool.name for tool in tools]))
        | llm.bind_functions(functions)
        | _update_state
    )
    workflow.add_node(name, chain)


def call_tool(state):
    messages = state["messages"]
    # Based on the continue condition
    # we know the last message involves a function call
    last_message = messages[-1]
    # We construct an ToolInvocation from the function_call
    action = ToolInvocation(
        tool=last_message.additional_kwargs["function_call"]["name"],
        tool_input=json.loads(
            last_message.additional_kwargs["function_call"]["arguments"]
        ),
    )
    # We call the tool_executor and get back a response
    response = tool_executor.invoke(action)
    # We use the response to create a FunctionMessage
    function_message = FunctionMessage(content=str(response), name=action.tool)
    # We return a list, because this will get added to the existing list
    return {"messages": [function_message]}


def should_continue(state):
    # This is the router
    messages = state["messages"]
    last_message = messages[-1]
    if "function_call" in last_message.additional_kwargs:
        return "call_tool"
    if "FINAL ANSWER" in last_message.content:
        return "end"
    return "continue"


llm = ChatOpenAI(model="gpt-4")
add_agent_node(
    "Researcher",
    llm,
    [tavily_tool],
    prompt,
)
add_agent_node("Chart Generator", llm, [create_plot], prompt)
workflow.add_node("call_tool", call_tool)
workflow.add_conditional_edges(
    "Researcher",
    should_continue,
    {"continue": "Chart Generator", "call_tool": "call_tool", "end": END},
)
workflow.add_conditional_edges(
    "Chart Generator",
    should_continue,
    {"continue": "Researcher", "call_tool": "call_tool", "end": END},
)
# We will let the user_info tool decide when to respond.
workflow.add_edge("call_tool", "Researcher")
workflow.set_entry_point("Researcher")
graph = workflow.compile()

Invoke

With the graph created, you can invoke it!

In [ ]:

result = graph.invoke(
    {
        "messages": [
            HumanMessage(
                content="Fetch the UK's GDP over the past 5 years, then draw a line graph of it."
            )
        ]
    },
    # Maximum number of steps to take in the graph
    {"recursion_limit": 150},
)
result["messages"][-1]