Files
langgraph/examples/multi_agent/agent_supervisor.ipynb
T

348 lines
14 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276",
"metadata": {},
"source": [
"## Agent Supervisor\n",
"\n",
"The [previous example](multi-agent-collaboration.ipynb) routed messages automatically based on the output of the initial researcher agent.\n",
"\n",
"We can also choose to use an LLM to orchestrate the different agents.\n",
"\n",
"Below, we will create an agent group, with an agent supervisor to help delegate tasks.\n",
"\n",
"![diagram](./img/supervisor-diagram.png)\n",
"\n",
"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.\n",
"\n",
"Before we build, let's configure our environment:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "0d30b6f7-3bec-4d9f-af50-43dfdc81ae6c",
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain langchain_openai langchain_experimental langsmith pandas"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "30c2f3de-c730-4aec-85a6-af2c2f058803",
"metadata": {},
"outputs": [],
"source": [
"import getpass\n",
"import os\n",
"\n",
"\n",
"def _set_if_undefined(var: str):\n",
" if not os.environ.get(var):\n",
" os.environ[var] = getpass.getpass(f\"Please provide your {var}\")\n",
"\n",
"\n",
"_set_if_undefined(\"OPENAI_API_KEY\")\n",
"_set_if_undefined(\"LANGCHAIN_API_KEY\")\n",
"_set_if_undefined(\"TAVILY_API_KEY\")\n",
"\n",
"# Optional, add tracing in LangSmith\n",
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
"os.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""
]
},
{
"cell_type": "markdown",
"id": "1ac25624-4d83-45a4-b9ef-a10589aacfb7",
"metadata": {},
"source": [
"## Create tools\n",
"\n",
"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:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "f04c6778-403b-4b49-9b93-678e910d5cec",
"metadata": {},
"outputs": [],
"source": [
"from typing import Annotated\n",
"\n",
"from langchain_community.tools.tavily_search import TavilySearchResults\n",
"from langchain_experimental.tools import PythonREPLTool\n",
"\n",
"tavily_tool = TavilySearchResults(max_results=5)\n",
"\n",
"# This executes code locally, which can be unsafe\n",
"python_repl_tool = PythonREPLTool()"
]
},
{
"cell_type": "markdown",
"id": "d58d1e85-22d4-4c22-9062-72a346a0d709",
"metadata": {},
"source": [
"## Helper Utilities"
]
},
{
"cell_type": "markdown",
"id": "b7c302b0-cd57-4913-986f-5dc7d6d77386",
"metadata": {},
"source": [
"Define a helper function that we will use to create the nodes in the graph - it takes care of converting the agent response to a human message. This is important because that is how we will add it the global state of the graph"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "80862241-a1a7-4726-bce5-f867b233832e",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.messages import HumanMessage\n",
"\n",
"def agent_node(state, agent, name):\n",
" result = agent.invoke(state)\n",
" return {\"messages\": [HumanMessage(content=result[\"messages\"][-1].content, name=name)]}"
]
},
{
"cell_type": "markdown",
"id": "d32962d2-5487-496d-aefc-2a3b0d194985",
"metadata": {},
"source": [
"### Create Agent Supervisor\n",
"\n",
"It will use function calling to choose the next worker node OR finish processing."
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "311f0a58-b425-4496-adac-dc4cd8ffb912",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
"from langchain_openai import ChatOpenAI\n",
"from pydantic import BaseModel\n",
"from typing import Literal\n",
"\n",
"members = [\"Researcher\", \"Coder\"]\n",
"system_prompt = (\n",
" \"You are a supervisor tasked with managing a conversation between the\"\n",
" \" following workers: {members}. Given the following user request,\"\n",
" \" respond with the worker to act next. Each worker will perform a\"\n",
" \" task and respond with their results and status. When finished,\"\n",
" \" respond with FINISH.\"\n",
")\n",
"# Our team supervisor is an LLM node. It just picks the next agent to process\n",
"# and decides when the work is completed\n",
"options = [\"FINISH\"] + members\n",
"\n",
"class routeResponse(BaseModel):\n",
" next: Literal[*options]\n",
"\n",
"prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
" (\"system\", system_prompt),\n",
" MessagesPlaceholder(variable_name=\"messages\"),\n",
" (\n",
" \"system\",\n",
" \"Given the conversation above, who should act next?\"\n",
" \" Or should we FINISH? Select one of: {options}\",\n",
" ),\n",
" ]\n",
").partial(options=str(options), members=\", \".join(members))\n",
"\n",
"\n",
"llm = ChatOpenAI(model=\"gpt-4o\")\n",
"\n",
"def supervisor_agent(state):\n",
" supervisor_chain = (\n",
" prompt\n",
" | llm.with_structured_output(routeResponse)\n",
" )\n",
" return supervisor_chain.invoke(state)"
]
},
{
"cell_type": "markdown",
"id": "a07d507f-34d1-4f1b-8dde-5e58d17b2166",
"metadata": {},
"source": [
"## Construct Graph\n",
"\n",
"We're ready to start building the graph. Below, define the state and worker nodes using the function we just defined."
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "6a430af7-8fce-4e66-ba9e-d940c1bc48e8",
"metadata": {},
"outputs": [],
"source": [
"import functools\n",
"import operator\n",
"from typing import Sequence, TypedDict\n",
"\n",
"from langchain_core.messages import BaseMessage\n",
"\n",
"from langgraph.graph import END, StateGraph, START\n",
"from langgraph.prebuilt import create_react_agent\n",
"\n",
"# The agent state is the input to each node in the graph\n",
"class AgentState(TypedDict):\n",
" # The annotation tells the graph that new messages will always\n",
" # be added to the current states\n",
" messages: Annotated[Sequence[BaseMessage], operator.add]\n",
" # The 'next' field indicates where to route to next\n",
" next: str\n",
"\n",
"\n",
"research_agent = create_react_agent(llm, tools=[tavily_tool])\n",
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n",
"\n",
"# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION\n",
"code_agent = create_react_agent(llm, tools=[python_repl_tool])\n",
"code_node = functools.partial(agent_node, agent=code_agent, name=\"Coder\")\n",
"\n",
"workflow = StateGraph(AgentState)\n",
"workflow.add_node(\"Researcher\", research_node)\n",
"workflow.add_node(\"Coder\", code_node)\n",
"workflow.add_node(\"supervisor\", supervisor_chain)"
]
},
{
"cell_type": "markdown",
"id": "2c1593d5-39f7-4819-96d2-4ad7d7991d72",
"metadata": {},
"source": [
"Now connect all the edges in the graph."
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "14778e86-077b-4e6a-893c-400e59b0cdbf",
"metadata": {},
"outputs": [],
"source": [
"for member in members:\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" workflow.add_edge(member, \"supervisor\")\n",
"# The supervisor populates the \"next\" field in the graph state\n",
"# which routes to a node or finishes\n",
"conditional_map = {k: k for k in members}\n",
"conditional_map[\"FINISH\"] = END\n",
"workflow.add_conditional_edges(\"supervisor\", lambda x: x[\"next\"], conditional_map)\n",
"# Finally, add entrypoint\n",
"workflow.add_edge(START, \"supervisor\")\n",
"\n",
"graph = workflow.compile()"
]
},
{
"cell_type": "markdown",
"id": "d36496de-7121-4c49-8cb6-58c943c66628",
"metadata": {},
"source": [
"## Invoke the team\n",
"\n",
"With the graph created, we can now invoke it and see how it performs!"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "56ba78e9-d9c1-457c-a073-d606d5d3e013",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'supervisor': {'next': 'Coder'}}\n",
"----\n",
"{'Coder': {'messages': [HumanMessage(content='The code to print \"Hello, World!\" to the terminal is:\\n\\n```python\\nprint(\\'Hello, World!\\')\\n```\\n\\nWhen executed, it prints:\\n```\\nHello, World!\\n```', name='Coder')]}}\n",
"----\n",
"{'supervisor': {'next': 'FINISH'}}\n",
"----\n"
]
}
],
"source": [
"for s in graph.stream(\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(content=\"Code hello world and print it to the terminal\")\n",
" ]\n",
" }\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"----\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "45a92dfd-0e11-47f5-aad4-b68d24990e34",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'supervisor': {'next': 'Researcher'}}\n",
"----\n",
"{'Researcher': {'messages': [HumanMessage(content='# Research Report on Pikas\\n\\nPikas, belonging to the genus Ochotona, are small, short-legged, and virtually tailless mammals that are often found in the mountains of western North America and across much of Asia. Despite their rodent-like appearance, pikas are not rodents but rather are part of the order Lagomorpha, which also includes rabbits and hares.\\n\\n## Behavior and Ecology\\nPikas are known for their unique behavior of not hibernating and remaining active throughout the winter. They navigate through tunnels under rocks and snow and rely on dried plants, which they have stored during warmer months in caches known as \"haypiles.\" This foraging strategy, termed \"haying,\" is crucial for their survival during the harsh winter months.\\n\\nPikas have a preference for cooler temperatures, typically foraging in temperatures below 25°C (77°F). They tend to avoid direct sunlight and stay in shaded regions when it gets warmer. A study has shown that for every 1°C (1.8°F) increase in ambient temperature, pikas can lose 3% of their foraging time, making them sensitive to climate change.\\n\\n## Distribution and Habitat\\nThe American pika (Ochotona princeps) and its relative, the collared pika (O. collaris), are found throughout the high mountainous regions of western North America. These species prefer cooler climates and have been observed to retreat to higher elevations as a response to increasing temperatures. Their current distribution is believed to be a result of a retreat from much larger ranges they occupied in the past, which included Western Europe and Eastern North America.\\n\\n## Conservation Status\\nThe International Union for Conservation of Nature and Natural Resources (IUCN) lists the American pika as a species of Least Concern but notes that populations are declining and unlikely to rebound due to habitat loss from extreme temperatures. The sensitivity of pikas to summer heat makes them an indicator species for the potential effects of climate change. Studies have shown that some populations are in decline, and there have been cases of local extirpation, particularly in the Great Basin.\\n\\n## Human Impact\\nHuman activity has impacted the ecosystems where pikas live, with recorded interactions dating back to the 1970s. Such interactions have been linked to pikas having reduced foraging time, limiting the amount of food they can stockpile for winter. Additionally, pikas have been considered pests in regions like the Tibetan plateau, where high densities of burrowing pikas are thought to reduce forage for domestic livestock and damage grasslands.\\n\\n## Conclusion\\nPikas are fascinating creatures with distinct adaptations that allow them to thrive in alpine environments. However, their future is uncertain due to the looming threats of climate change and habitat alteration. Conservation efforts, research, and monitoring are vital to ensure the survival of these unique mammals in a changing world.\\n\\n---\\n\\n**Sources:**\\n- [Wikipedia - Pika](https://en.wikipedia.org/wiki/Pika)\\n- [Treehugger - American Pika](https://www.treehugger.com/surprising-facts-about-american-pika-4864528)\\n- [National Park Service - Pikas at Rocky Mountain National Park](https://www.nps.gov/romo/learn/nature/pikas.htm)\\n- [Wikipedia - American Pika](https://en.wikipedia.org/wiki/American_pika)\\n- [Britannica - Pika](https://www.britannica.com/animal/pika)', name='Researcher')]}}\n",
"----\n",
"{'supervisor': {'next': 'FINISH'}}\n",
"----\n"
]
}
],
"source": [
"for s in graph.stream(\n",
" {\"messages\": [HumanMessage(content=\"Write a brief research report on pikas.\")]},\n",
" {\"recursion_limit\": 100},\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"----\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}