Files
langgraph/examples/multi_agent/agent_supervisor.ipynb
T
2024-01-20 14:20:24 -08:00

16 KiB
Raw Blame History

Agent Supervisor

The previous example routed messages automatically based on the output of the initial researcher agent.

We can also choose to use an LLM to orchestrate the different agents.

Below, we will create an agent group, with an agent supervisor to help delegate tasks.

diagram

To simplify the code in each agent node, we will use the AgentExecutor class from LangChain. 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.

Before we build, let's configure our environment:

In [1]:
# %%capture --no-stderr
# %pip install -U langchain langchain_openai langchain_experimental langsmith pandas
In [2]:
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 tools

For this example, you will make an agent to do web research with a search engine, and one agent to create plots. Define the tools they'll use below:

In [3]:
from typing import Annotated, List, Tuple, Union

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool
from langchain_experimental.tools import PythonREPLTool

tavily_tool = TavilySearchResults(max_results=5)

# This executes code locally, which can be unsafe
python_repl_tool = PythonREPLTool()

Helper Utilites

Define a helper function below, which make it easier to add new agent worker nodes.

In [4]:
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI

from langgraph.graph import END, StateGraph


def create_worker_node(
    workflow: StateGraph, name: str, llm: ChatOpenAI, tools: list, system_prompt: str
):
    # Each worker node will be given a name and some tools.
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                system_prompt,
            ),
            MessagesPlaceholder(variable_name="messages"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    agent = create_openai_tools_agent(llm, tools, prompt)
    executor = AgentExecutor(agent=agent, tools=tools)
    chain = executor | (
        # So the agents properly role-play in this simulation, we will
        # tag their final message as a human message
        lambda x: {"messages": [HumanMessage(content=x["output"], name=name)]}
    )
    workflow.add_node(name, chain)

Construct Graph

We're ready to start building the graph. Below, define the state and worker nodes using the function we just defined.

In [5]:
import operator
from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder


# The agent state is the input to each node in the graph
class AgentState(TypedDict):
    # The annotation tells the graph that new messages will always
    # be added to the current states
    messages: Annotated[Sequence[BaseMessage], operator.add]
    # The 'next' field indicates where to route to next
    next: str


workflow = StateGraph(AgentState)

llm = ChatOpenAI(model="gpt-4-1106-preview")


create_worker_node(
    workflow, "Researcher", llm, [tavily_tool], "You are a web researcher."
)
# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION
create_worker_node(
    workflow,
    "Coder",
    llm,
    [python_repl_tool],
    "You may generate safe python code to analyze data "
    "and generate charts using matplotlib.",
)

Almost done, now create create the team supervisor. It will use function calling to choose the next worker node OR finish processing.

In [6]:
from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser

members = ["Researcher", "Coder"]
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. When finished,"
    " respond with FINISH."
)
# Our team supervisor is an LLM node. It just picks the next agent to process
# and decides when the work is completed
options = ["FINISH"] + members
# Using openai function calling can make output parsing easier for us
function_def = {
    "name": "route",
    "description": "Select the next role.",
    "parameters": {
        "title": "routeSchema",
        "type": "object",
        "properties": {
            "next": {
                "title": "Next",
                "anyOf": [
                    {"enum": options},
                ],
            }
        },
        "required": ["next"],
    },
}
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="messages"),
        (
            "system",
            "Given the conversation above, who should act next?"
            " Or should we FINISH? Select one of: {options}",
        ),
    ]
).partial(options=str(options), members=", ".join(members))

supervisor_chain = (
    prompt
    | llm.bind_functions(functions=[function_def], function_call="route")
    | JsonOutputFunctionsParser()
)

Now connect all the edges in the graph.

In [7]:
workflow.add_node("supervisor", supervisor_chain)


for member in members:
    # We want our workers to ALWAYS "report back" to the supervisor when done
    workflow.add_edge(member, "supervisor")
# The supervisor populates the "next" field in the graph state
# which routes to a node or finishes
conditional_map = {k: k for k in members}
conditional_map["FINISH"] = END
workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
# Finally, add entrypoint
workflow.set_entry_point("supervisor")

graph = workflow.compile()

Invoke the team

With the graph created, we can now invoke it and see how it performs!

In [8]:
results = graph.invoke(
    {
        "messages": [
            HumanMessage(content="Code hello world and print it to the terminal")
        ]
    }
)
results["messages"][-1].pretty_print()
Python REPL can execute arbitrary code. Use with caution.
================================ Human Message =================================

The code `print('Hello, World!')` was executed, and the output is:

```
Hello, World!
```
In [12]:
results = graph.invoke(
    {
        "messages": [
            HumanMessage(content="Write a brief research report on pikas.")
        ]
    },
    {"recursion_limit": 100},
)
results["messages"][-1].pretty_print()
================================ Human Message =================================

# Research Report on Pikas

Pikas are small, mountain-dwelling mammals that are closely related to rabbits. They are known for their distinctive chirps and typically inhabit boulder fields at high elevations, up to 14,000 feet in treeless slopes like those found in the Southern Rockies. These animals are recognized for their ability to adapt to some of the most inhospitable climates.

## Climate Change Impact

Pikas have been a topic of interest in climate change research due to their sensitivity to high temperatures and reliance on cold habitats. They have historically responded to climate shifts by moving to higher elevations or latitudes to find suitable cooler environments. For instance, pikas were once found in the Appalachian Mountains and even in the Mojave Desert, but as the Earth's climate warmed, they moved to cooler, high-elevation areas where they live today.

Recent studies suggest that pikas are showing remarkable adaptability to climate change. Despite predictions that they might become endangered due to rising temperatures, these animals are displaying resilience. Some research indicates that pikas can adjust certain genes to make better use of oxygen in higher altitudes where the air is thinner, which could be a potential hope for their survival as climate change drives them to higher elevations.

However, there have been reports of pikas disappearing from parts of the Great Basin, and in Colorado, pikas have retracted upslope by about 1,160 feet. It's been noted that while climate change may be a factor, it might not be the sole cause for these local disappearances.

## Adaptation Strategies

Pikas exhibit several interesting behaviors that help them cope with their challenging environment. During the summer, they engage in activities like "making hay" — collecting and storing vegetation in preparation for the harsh winters. Their diet and caching behavior are essential for their survival during the months when food is scarce.

## Conservation and Research

Conservationists and scientists continue to study pikas to understand their adaptation mechanisms and how they might inform broader climate change mitigation strategies. For example, studies have been conducted on pikas at different elevations to observe genetic changes and their effects on adaptation. Such research is crucial for predicting the future of pikas and potentially other species affected by climate change.

## Conclusion

Pikas serve as an important indicator species for the impacts of climate change on wildlife. Their ability to adapt to changing climates offers hope and also highlights the importance of understanding genetic adaptability in the face of environmental challenges. Conservation efforts and further research are essential to ensure the survival of pikas and to learn from their resilience.

### Sources
- [The Conversation: Pikas are adapting to climate change remarkably well](https://theconversation.com/pikas-are-adapting-to-climate-change-remarkably-well-contrary-to-many-predictions-150726)
- [Stanford Sustainability: It's in the genes  potential hope for pikas hit by climate change](https://sustainability.stanford.edu/news/its-genes-potential-hope-pikas-hit-climate-change)
- [Colorado Sun: Colorado pika population and climate change](https://coloradosun.com/2023/08/27/colorado-pika-population-climate-change/)
- [PetaPixel: Photographing the American Pika  a tiny indicator of climate change](https://petapixel.com/2022/01/03/photographing-the-american-pika-a-tiny-indicator-of-climate-change/)
- [The Wildlife Society: Can pikas survive climate change after all?](https://wildlife.org/can-pikas-survive-climate-change-after-all/)
In [ ]: