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

Multi-agent Collaboration Intro

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.

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"

Example 1: 2 Agents

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

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

In [ ]:

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()

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]