langgraph/examples/multi_agent/agent_supervisor.ipynb

{
 "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 langchain langchain_openai langchain_experimental langsmith pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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(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, List, Tuple, Union\n",
    "\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_core.tools import tool\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 Utilites\n",
    "\n",
    "Define a helper function below, which make it easier to add new agent worker nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c4823dd9-26bd-4e1a-8117-b97b2860211a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, create_openai_tools_agent\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "\n",
    "def create_agent(\n",
    "    llm: ChatOpenAI, tools: list, system_prompt: str\n",
    "):\n",
    "    # Each worker node will be given a name and some tools.\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\n",
    "                \"system\",\n",
    "                system_prompt,\n",
    "            ),\n",
    "            MessagesPlaceholder(variable_name=\"messages\"),\n",
    "            MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
    "        ]\n",
    "    )\n",
    "    agent = create_openai_tools_agent(llm, tools, prompt)\n",
    "    executor = AgentExecutor(agent=agent, tools=tools)\n",
    "    chain = executor | (\n",
    "        # So the agents properly role-play in this simulation, we will\n",
    "        # tag their final message as a human message\n",
    "        lambda x: {\"messages\": [HumanMessage(content=x[\"output\"], name=name)]}\n",
    "    )\n",
    "    workflow.add_node(name, chain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c302b0-cd57-4913-986f-5dc7d6d77386",
   "metadata": {},
   "source": [
    "We can also define a function that we will use to be 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": 3,
   "id": "80862241-a1a7-4726-bce5-f867b233832e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def agent_node(state, agent, name):\n",
    "    result = agent.invoke(state)\n",
    "    return {\"messages\": [HumanMessage(content=result[\"output\"], 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": 6,
   "id": "311f0a58-b425-4496-adac-dc4cd8ffb912",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser\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",
    "# Using openai function calling can make output parsing easier for us\n",
    "function_def = {\n",
    "    \"name\": \"route\",\n",
    "    \"description\": \"Select the next role.\",\n",
    "    \"parameters\": {\n",
    "        \"title\": \"routeSchema\",\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"next\": {\n",
    "                \"title\": \"Next\",\n",
    "                \"anyOf\": [\n",
    "                    {\"enum\": options},\n",
    "                ],\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"next\"],\n",
    "    },\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",
    "supervisor_chain = (\n",
    "    prompt\n",
    "    | llm.bind_functions(functions=[function_def], function_call=\"route\")\n",
    "    | JsonOutputFunctionsParser()\n",
    ")"
   ]
  },
  {
   "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": 12,
   "id": "6a430af7-8fce-4e66-ba9e-d940c1bc48e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict\n",
    "import functools\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langgraph.graph import StateGraph, END\n",
    "\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",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
    "\n",
    "research_agent = create_agent(llm, [tavily_tool], \"You are a web researcher.\")\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_agent(llm, [python_repl_tool], \"You may generate safe python code to analyze data and generate charts using matplotlib.\")\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": 13,
   "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.set_entry_point(\"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": 14,
   "id": "56ba78e9-d9c1-457c-a073-d606d5d3e013",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'supervisor': {'next': 'Coder'}}\n",
      "----\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Python REPL can execute arbitrary code. Use with caution.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Coder': {'messages': [HumanMessage(content=\"The code `print('Hello, World!')` was executed, and it printed `Hello, World!` to the terminal.\", 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": 16,
   "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 are small, mountain-dwelling mammals native to Asia and North America. They are characterized by their short limbs, rounded ears, and stout bodies. Unlike rodents, with which they are sometimes confused, pikas are more closely related to rabbits and hares, belonging to the order Lagomorpha.\\n\\n### Physical Description\\nPikas are typically 15 to 23 cm (5.9 to 9.1 inches) in body length. They have an even coat of fur, no external tail, and resemble their close relatives, the rabbits, but with short, rounded ears. Their thick fur and small ears are adaptations to their cold habitats.\\n\\n### Habitat and Distribution\\nPikas inhabit mountainous regions where they have adapted to life in rocky terrains. These creatures can be found at high altitudes, with the large-eared pika of the Himalayas living at elevations of more than 6,000 meters.\\n\\n### Diet and Behavior\\nPikas are herbivorous and primarily feed on plants. They are known for their diligent behavior of collecting and storing food in their tunnels for winter. During the warmer months, pikas spend their days gathering grass in montane meadows and creating hay piles. These animals are also remarkable for their distinctive vocalizations used for communication.\\n\\n### Reproduction and Lifespan\\nDetails on pika reproduction and lifespan were not provided in the immediate search results, but like many small mammals, pikas would typically have a breeding season and a relatively short lifespan in the wild.\\n\\n### Conservation\\nPikas are currently facing challenges due to environmental changes. As their habitats are affected by the expansion of agriculture and the increasing temperatures due to climate change, pikas may find themselves in competition for resources. They are notably heat intolerant, which makes them particularly vulnerable to global warming.\\n\\n### Conclusion\\nPikas are fascinating creatures with unique adaptations to their high-altitude environments. Their role in the ecosystem, distinctive behaviors, and close relation to the rabbit make them an interesting subject for further study, especially in the context of environmental conservation.\\n\\n**Sources:**\\n- [Animals.net - Pika](https://animals.net/pika/)\\n- [OneKindPlanet - American Pika](https://www.onekindplanet.org/animal/american-pika/)\\n- [Facts.net - 12 Facts About Pika](https://facts.net/nature/animals/12-facts-about-pika/)\\n- [Wikipedia - Pika](https://en.wikipedia.org/wiki/Pika)\\n- [HowStuffWorks - Pikas](https://animals.howstuffworks.com/mammals/pika.htm)', name='Researcher')]}}\n",
      "----\n",
      "{'supervisor': {'next': 'FINISH'}}\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "for s in graph.stream(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            HumanMessage(content=\"Write a brief research report on pikas.\")\n",
    "        ]\n",
    "    },\n",
    "    {\"recursion_limit\": 100},\n",
    "):\n",
    "    if \"__end__\" not in s:\n",
    "        print(s)\n",
    "        print(\"----\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d363d2c-e0da-4cce-ba47-ad2aa9df0fef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}