examples/multi_agent/hierarchical_agent_teams.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276",
   "metadata": {},
   "source": [
    "## Hierarchical Agent Teams\n",
    "\n",
    "In our previous example ([Agent Supervisor](./agent_supervisor.ipynb)), we introduced the concept of a single supervisor node to route work between different worker nodes.\n",
    "\n",
    "But what if the job for a single worker becomes too complex? What if the number of workers becomes too large?\n",
    "\n",
    "For some applications, the system may be more effective if work is distributed _hierarchically_.\n",
    "\n",
    "You can do this by composing different subgraphs and creating a top-level supervisor, along with mid-level supervisors.\n",
    "\n",
    "To do this, let's build a simple research assistant! The graph will look something like the following:\n",
    "\n",
    "![diagram](./img/hierarchical-diagram.png)\n",
    "\n",
    "This notebook is inspired by the paper [AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation](https://arxiv.org/abs/2308.08155), by Wu, et. al. In the rest of this notebook, you will:\n",
    "\n",
    "1. Define the agents' tools to access the web and write files\n",
    "2. Define some utilities to help create the graph and agents\n",
    "3. Create and define each team (web research + doc writing)\n",
    "4. Compose everything together.\n",
    "\n",
    "But before all of that, some setup:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d30b6f7-3bec-4d9f-af50-43dfdc81ae6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%capture --no-stderr\n",
    "# %pip install -U langgraph langchain langchain_openai langchain_experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30c2f3de-c730-4aec-85a6-af2c2f058803",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "import uuid\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",
    "# This will help you visualize and debug the control flow\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "354568e2-aef0-4af9-8a79-e64d3eea752f",
   "metadata": {},
   "source": [
    "## Create Tools\n",
    "\n",
    "Each team will be composed of one or more agents each with one or more tools. Below, define all the tools to be used by your different teams.\n",
    "\n",
    "We'll start with the research team.\n",
    "\n",
    "**Research team tools**\n",
    "\n",
    "The research team can use a search engine and url scraper to find information on the web. Feel free to add additional functionality below to boost the team performance!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4024eb89-843d-4cc3-ab3f-e1eb4d031179",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List, Tuple, Union\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_core.tools import tool\n",
    "from langsmith import trace\n",
    "\n",
    "tavily_tool = TavilySearchResults(max_results=5)\n",
    "\n",
    "\n",
    "@tool\n",
    "def scrape_webpages(urls: List[str]) -> str:\n",
    "    \"\"\"Use requests and bs4 to scrape the provided web pages for detailed information.\"\"\"\n",
    "    loader = WebBaseLoader(urls)\n",
    "    docs = loader.load()\n",
    "    return \"\\n\\n\".join(\n",
    "        [\n",
    "            f'<Document name=\"{doc.metadata.get(\"title\", \"\")}\">\\n{doc.page_content}\\n</Document>'\n",
    "            for doc in docs\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c427982-fadf-4721-a77e-2465df9fc6bc",
   "metadata": {},
   "source": [
    "**Document writing team tools**\n",
    "\n",
    "Next up, we will give some tools for the doc writing team to use.\n",
    "We define some bare-bones file-access tools below.\n",
    "\n",
    "Note that this gives the agents access to your file-system, which can be unsafe. We also haven't optimized the tool descriptions for performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f20a18ca-2709-4c12-84f3-88678591a9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from tempfile import TemporaryDirectory\n",
    "from typing import Dict, Optional\n",
    "\n",
    "from langchain_experimental.utilities import PythonREPL\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "_TEMP_DIRECTORY = TemporaryDirectory()\n",
    "WORKING_DIRECTORY = Path(_TEMP_DIRECTORY.name)\n",
    "\n",
    "\n",
    "@tool\n",
    "def create_outline(\n",
    "    points: Annotated[List[str], \"List of main points or sections.\"],\n",
    "    file_name: Annotated[str, \"File path to save the outline.\"],\n",
    ") -> Annotated[str, \"Path of the saved outline file.\"]:\n",
    "    \"\"\"Create and save an outline.\"\"\"\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        for i, point in enumerate(points):\n",
    "            file.write(f\"{i + 1}. {point}\\n\")\n",
    "    return f\"Outline saved to {file_name}\"\n",
    "\n",
    "\n",
    "@tool\n",
    "def read_document(\n",
    "    file_name: Annotated[str, \"File path to save the document.\"],\n",
    "    start: Annotated[Optional[int], \"The start line. Default is 0\"] = None,\n",
    "    end: Annotated[Optional[int], \"The end line. Default is None\"] = None,\n",
    ") -> str:\n",
    "    \"\"\"Read the specified document.\"\"\"\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"r\") as file:\n",
    "        lines = file.readlines()\n",
    "    if start is not None:\n",
    "        start = 0\n",
    "    return \"\\n\".join(lines[start:end])\n",
    "\n",
    "\n",
    "@tool\n",
    "def write_document(\n",
    "    content: Annotated[str, \"Text content to be written into the document.\"],\n",
    "    file_name: Annotated[str, \"File path to save the document.\"],\n",
    ") -> Annotated[str, \"Path of the saved document file.\"]:\n",
    "    \"\"\"Create and save a text document.\"\"\"\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        file.write(content)\n",
    "    return f\"Document saved to {file_name}\"\n",
    "\n",
    "\n",
    "@tool\n",
    "def edit_document(\n",
    "    file_name: Annotated[str, \"Path of the document to be edited.\"],\n",
    "    inserts: Annotated[\n",
    "        Dict[int, str],\n",
    "        \"Dictionary where key is the line number (1-indexed) and value is the text to be inserted at that line.\",\n",
    "    ],\n",
    ") -> Annotated[str, \"Path of the edited document file.\"]:\n",
    "    \"\"\"Edit a document by inserting text at specific line numbers.\"\"\"\n",
    "\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"r\") as file:\n",
    "        lines = file.readlines()\n",
    "\n",
    "    sorted_inserts = sorted(inserts.items())\n",
    "\n",
    "    for line_number, text in sorted_inserts:\n",
    "        if 1 <= line_number <= len(lines) + 1:\n",
    "            lines.insert(line_number - 1, text + \"\\n\")\n",
    "        else:\n",
    "            return f\"Error: Line number {line_number} is out of range.\"\n",
    "\n",
    "    with (WORKING_DIRECTORY / file_name).open(\"w\") as file:\n",
    "        file.writelines(lines)\n",
    "\n",
    "    return f\"Document edited and saved to {file_name}\"\n",
    "\n",
    "\n",
    "# Warning: This executes code locally, which can be unsafe when not sandboxed\n",
    "\n",
    "repl = PythonREPL()\n",
    "\n",
    "\n",
    "@tool\n",
    "def python_repl(\n",
    "    code: Annotated[str, \"The python code to execute to generate your chart.\"]\n",
    "):\n",
    "    \"\"\"Use this to execute python code. If you want to see the output of a value,\n",
    "    you should print it out with `print(...)`. This is visible to the user.\"\"\"\n",
    "    try:\n",
    "        result = repl.run(code)\n",
    "    except BaseException as e:\n",
    "        return f\"Failed to execute. Error: {repr(e)}\"\n",
    "    return f\"Successfully executed:\\n```python\\n{code}\\n```\\nStdout: {result}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504ee1c6-2b6a-439d-9046-df54e1e15698",
   "metadata": {},
   "source": [
    "## Helper Utilities\n",
    "\n",
    "We are going to create a few utility functions to make it more concise when we want to:\n",
    "\n",
    "1. Create a worker agent.\n",
    "2. Create a supervisor for the sub-graph.\n",
    "\n",
    "These will simplify the graph compositional code at the end for us so it's easier to see what's going on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e09fb60f-1aac-455b-b67d-8d2e4ccfd747",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Callable, List, Optional, TypedDict, Union\n",
    "\n",
    "from langchain.agents import AgentExecutor, create_openai_functions_agent\n",
    "from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.runnables import Runnable\n",
    "from langchain_core.tools import BaseTool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph.graph import END, StateGraph\n",
    "\n",
    "\n",
    "def create_agent(\n",
    "    llm: ChatOpenAI,\n",
    "    tools: list,\n",
    "    system_prompt: str,\n",
    ") -> str:\n",
    "    \"\"\"Create a function-calling agent and add it to the graph.\"\"\"\n",
    "    system_prompt += \"\\nWork autonomously according to your specialty, using the tools available to you.\"\n",
    "    \" Do not ask for clarification.\"\n",
    "    \" Your other team members (and other teams) will collaborate with you with their own specialties.\"\n",
    "    \" You are chosen for a reason! You are one of the following team members: {team_members}.\"\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_functions_agent(llm, tools, prompt)\n",
    "    executor = AgentExecutor(agent=agent, tools=tools)\n",
    "    return executor\n",
    "\n",
    "\n",
    "def agent_node(state, agent, name):\n",
    "    result = agent.invoke(state)\n",
    "    return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}\n",
    "\n",
    "\n",
    "def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> str:\n",
    "    \"\"\"An LLM-based router.\"\"\"\n",
    "    options = [\"FINISH\"] + members\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), team_members=\", \".join(members))\n",
    "    return (\n",
    "        prompt\n",
    "        | llm.bind_functions(functions=[function_def], function_call=\"route\")\n",
    "        | JsonOutputFunctionsParser()\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00282b1f-bb4d-4ee7-9bae-e8e6f586f12e",
   "metadata": {},
   "source": [
    "## Define Agent Teams\n",
    "\n",
    "Now we can get to define our hierarchical teams. \"Choose your player!\"\n",
    "\n",
    "### Research Team\n",
    "\n",
    "The research team will have a search agent and a web scraping \"research_agent\" as the two worker nodes. Let's create those, as well as the team supervisor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "53db0c78-e357-48ba-ae5f-3fc04735a3b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "import operator\n",
    "\n",
    "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "import functools\n",
    "\n",
    "\n",
    "# Research team graph state\n",
    "class ResearchTeamState(TypedDict):\n",
    "    # A message is added after each team member finishes\n",
    "    messages: Annotated[List[BaseMessage], operator.add]\n",
    "    # The team members are tracked so they are aware of\n",
    "    # the others' skill-sets\n",
    "    team_members: List[str]\n",
    "    # Used to route work. The supervisor calls a function\n",
    "    # that will update this every time it makes a decision\n",
    "    next: str\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
    "\n",
    "search_agent = create_agent(\n",
    "    llm,\n",
    "    [tavily_tool],\n",
    "    \"You are a research assistant who can search for up-to-date info using the tavily search engine.\",\n",
    ")\n",
    "search_node = functools.partial(agent_node, agent=search_agent, name=\"Search\")\n",
    "\n",
    "research_agent = create_agent(\n",
    "    llm,\n",
    "    [scrape_webpages],\n",
    "    \"You are a research assistant who can scrape specified urls for more detailed information using the scrape_webpages function.\",\n",
    ")\n",
    "research_node = functools.partial(agent_node, agent=research_agent, name=\"Web Scraper\")\n",
    "\n",
    "supervisor_agent = create_team_supervisor(\n",
    "    llm,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following workers:  Search, Web Scraper. 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",
    "    [\"Search\", \"Web Scraper\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b01c6ee8-a461-4081-8a97-a3a06ec0f994",
   "metadata": {},
   "source": [
    "Now that we've created the necessary components, defining their interactions is easy. Add the nodes to the team graph, and define the edges, which determine the transition criteria."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1a7a1260-d9f6-4011-b2b1-13fab5126997",
   "metadata": {},
   "outputs": [],
   "source": [
    "research_graph = StateGraph(ResearchTeamState)\n",
    "research_graph.add_node(\"Search\", search_node)\n",
    "research_graph.add_node(\"Web Scraper\", research_node)\n",
    "research_graph.add_node(\"supervisor\", supervisor_agent)\n",
    "\n",
    "# Define the control flow\n",
    "research_graph.add_edge(\"Search\", \"supervisor\")\n",
    "research_graph.add_edge(\"Web Scraper\", \"supervisor\")\n",
    "research_graph.add_conditional_edges(\n",
    "    \"supervisor\",\n",
    "    lambda x: x[\"next\"],\n",
    "    {\"Search\": \"Search\", \"Web Scraper\": \"Web Scraper\", \"FINISH\": END},\n",
    ")\n",
    "\n",
    "\n",
    "research_graph.set_entry_point(\"supervisor\")\n",
    "chain = research_graph.compile()\n",
    "\n",
    "\n",
    "# The following functions interoperate between the top level graph state\n",
    "# and the state of the research sub-graph\n",
    "# this makes it so that the states of each graph don't get intermixed\n",
    "def enter_chain(message: str):\n",
    "    results = {\n",
    "        \"messages\": [HumanMessage(content=message)],\n",
    "    }\n",
    "    return results\n",
    "\n",
    "\n",
    "research_chain = enter_chain | chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ee8f2c-fbde-427b-ba54-ae0c7ce5fbfb",
   "metadata": {},
   "source": [
    "We can give this team work directly. Try it out below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "912b0604-a178-4246-a36f-2dedae606680",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'supervisor': {'next': 'Search'}}\n",
      "---\n",
      "{'Search': {'messages': [HumanMessage(content='Taylor Swift\\'s next tour, named \"The Eras Tour,\" has dates scheduled for both 2023 and 2024. Some of the upcoming dates mentioned in the search results include:\\n\\n- November 25-26, 2023, in São Paulo, Brazil at the Allianz Parque\\n- February 7-9, 2024, in Tokyo, Japan at the Tokyo Dome\\n\\nFor a complete and detailed list of all the dates and locations for Taylor Swift\\'s \"The Eras Tour,\" you would need to visit official sources or ticketing websites such as Ticketmaster, as the search results suggest additional dates have been added due to overwhelming demand. The tour is set to wrap up in Indianapolis on November 3, 2024. Keep in mind that the schedule may be subject to change, and it\\'s a good idea to verify the dates on an official ticketing or event site for the latest information.', name='Search')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'FINISH'}}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for s in research_chain.stream(\n",
    "    \"when is Taylor Swift's next tour?\", {\"recursion_limit\": 100}\n",
    "):\n",
    "    if \"__end__\" not in s:\n",
    "        print(s)\n",
    "        print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749b99ab-f6f0-4c5d-a90b-10102465d186",
   "metadata": {},
   "source": [
    "### Document Writing Team\n",
    "\n",
    "Create the document writing team below using a similar approach. This time, we will give each agent access to different file-writing tools.\n",
    "\n",
    "Note that we are giving file-system access to our agent here, which is not safe in all cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1bcdbf44-9481-430c-8429-fa142ed8a626",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "# Document writing team graph state\n",
    "class DocWritingState(TypedDict):\n",
    "    # This tracks the team's conversation internally\n",
    "    messages: Annotated[List[BaseMessage], operator.add]\n",
    "    # This provides each worker with context on the others' skill sets\n",
    "    team_members: str\n",
    "    # This is how the supervisor tells langgraph who to work next\n",
    "    next: str\n",
    "    # This tracks the shared directory state\n",
    "    current_files: str\n",
    "\n",
    "\n",
    "# This will be run before each worker agent begins work\n",
    "# It makes it so they are more aware of the current state\n",
    "# of the working directory.\n",
    "def prelude(state):\n",
    "    written_files = []\n",
    "    if not WORKING_DIRECTORY.exists():\n",
    "        WORKING_DIRECTORY.mkdir()\n",
    "    try:\n",
    "        written_files = [\n",
    "            f.relative_to(WORKING_DIRECTORY) for f in WORKING_DIRECTORY.rglob(\"*\")\n",
    "        ]\n",
    "    except:\n",
    "        pass\n",
    "    if not written_files:\n",
    "        return {**state, \"current_files\": \"No files written.\"}\n",
    "    return {\n",
    "        **state,\n",
    "        \"current_files\": \"\\nBelow are files your team has written to the directory:\\n\"\n",
    "        + \"\\n\".join([f\" - {f}\" for f in written_files]),\n",
    "    }\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
    "\n",
    "doc_writer_agent = create_agent(\n",
    "    llm,\n",
    "    [write_document, edit_document, read_document],\n",
    "    \"You are an expert writing a research document.\\n\"\n",
    "    # The {current_files} value is populated automatically by the graph state\n",
    "    \"Below are files currently in your directory:\\n{current_files}\",\n",
    ")\n",
    "# Injects current directory working state before each call\n",
    "context_aware_doc_writer_agent = prelude | doc_writer_agent\n",
    "doc_writing_node = functools.partial(\n",
    "    agent_node, agent=context_aware_doc_writer_agent, name=\"DocWriter\"\n",
    ")\n",
    "\n",
    "note_taking_agent = create_agent(\n",
    "    llm,\n",
    "    [create_outline, read_document],\n",
    "    \"You are an expert senior researcher tasked with writing a paper outline and\"\n",
    "    \" taking notes to craft a perfect paper.{current_files}\",\n",
    ")\n",
    "context_aware_note_taking_agent = prelude | note_taking_agent\n",
    "note_taking_node = functools.partial(\n",
    "    agent_node, agent=context_aware_note_taking_agent, name=\"NoteTaker\"\n",
    ")\n",
    "\n",
    "chart_generating_agent = create_agent(\n",
    "    llm,\n",
    "    [read_document, python_repl],\n",
    "    \"You are a data viz expert tasked with generating charts for a research project.\"\n",
    "    \"{current_files}\",\n",
    ")\n",
    "context_aware_chart_generating_agent = prelude | chart_generating_agent\n",
    "chart_generating_node = functools.partial(\n",
    "    agent_node, agent=context_aware_note_taking_agent, name=\"ChartGenerator\"\n",
    ")\n",
    "\n",
    "doc_writing_supervisor = create_team_supervisor(\n",
    "    llm,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following workers:  {team_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",
    "    [\"DocWriter\", \"NoteTaker\", \"ChartGenerator\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee2cd9b-29aa-458e-903d-4e49179e5d59",
   "metadata": {},
   "source": [
    "With the objects themselves created, we can form the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9c5c644f-8966-4d2e-98d2-80d73520e9fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the graph here:\n",
    "# Note that we have unrolled the loop for the sake of this doc\n",
    "authoring_graph = StateGraph(DocWritingState)\n",
    "authoring_graph.add_node(\"DocWriter\", doc_writing_node)\n",
    "authoring_graph.add_node(\"NoteTaker\", note_taking_node)\n",
    "authoring_graph.add_node(\"ChartGenerator\", chart_generating_node)\n",
    "authoring_graph.add_node(\"supervisor\", doc_writing_supervisor)\n",
    "\n",
    "# Add the edges that always occur\n",
    "authoring_graph.add_edge(\"DocWriter\", \"supervisor\")\n",
    "authoring_graph.add_edge(\"NoteTaker\", \"supervisor\")\n",
    "authoring_graph.add_edge(\"ChartGenerator\", \"supervisor\")\n",
    "\n",
    "# Add the edges where routing applies\n",
    "authoring_graph.add_conditional_edges(\n",
    "    \"supervisor\",\n",
    "    lambda x: x[\"next\"],\n",
    "    {\n",
    "        \"DocWriter\": \"DocWriter\",\n",
    "        \"NoteTaker\": \"NoteTaker\",\n",
    "        \"ChartGenerator\": \"ChartGenerator\",\n",
    "        \"FINISH\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "authoring_graph.set_entry_point(\"supervisor\")\n",
    "chain = authoring_graph.compile()\n",
    "\n",
    "\n",
    "# The following functions interoperate between the top level graph state\n",
    "# and the state of the research sub-graph\n",
    "# this makes it so that the states of each graph don't get intermixed\n",
    "def enter_chain(message: str, members: List[str]):\n",
    "    results = {\n",
    "        \"messages\": [HumanMessage(content=message)],\n",
    "        \"team_members\": \", \".join(members),\n",
    "    }\n",
    "    return results\n",
    "\n",
    "\n",
    "# We reuse the enter/exit functions to wrap the graph\n",
    "authoring_chain = (\n",
    "    functools.partial(enter_chain, members=authoring_graph.nodes)\n",
    "    | authoring_graph.compile()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9860fd46-c24d-40a5-a6ba-e8fddcd43369",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'supervisor': {'next': 'NoteTaker'}}\n",
      "---\n",
      "{'NoteTaker': {'messages': [HumanMessage(content='The outline and poem have been successfully written and saved. The poem, titled \"Whispers of the Wind,\" captures the essence of the wind\\'s journey from a gentle beginning to a powerful crescendo and then a calming end, concluding with the cyclical nature of this elemental force.\\n\\nHere is the completed poem:\\n\\n---\\n\\n**Whispers of the Wind**\\n\\nEmbracing the essence of nature\\'s breath,\\nA zephyr stirs, the dawn\\'s own cheer,\\nWhispering secrets, meant for no ear,\\nThe gentle beginnings, a soft caress.\\n\\nA gale ascends, midday\\'s fierce guest,\\nTwirling leaves in a wild, untamed dance,\\nNature\\'s raw power in its vast expanse,\\nThe crescendo of gusts, a forceful dance.\\n\\nAs dusk falls, the tempest does wane,\\nA soft sighing through the willow\\'s mane,\\nIn the quietude, peace finds its chance,\\nThe calming aftermath, a serene trance.\\n\\nThrough whispers, roars, and silent chants,\\nThe wind carries on, no pause nor glance,\\nIn its journey, a boundless, spirited lance,\\nThe enduring cycle, a timeless romance.\\n\\n---\\n\\nThe poem is a testament to the wind\\'s ever-present and ever-changing presence in the world around us.', name='NoteTaker')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'FINISH'}}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for s in authoring_chain.stream(\n",
    "    \"Write an outline for poem and then write the poem to disk.\",\n",
    "    {\"recursion_limit\": 100},\n",
    "):\n",
    "    if \"__end__\" not in s:\n",
    "        print(s)\n",
    "        print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4b5b08d-9a9a-474a-94b4-f7aaa8ff19e6",
   "metadata": {},
   "source": [
    "## Add Layers\n",
    "\n",
    "In this design, we are enforcing a top-down planning policy. We've created two graphs already, but we have to decide how to route work between the two.\n",
    "\n",
    "We'll create a _third_ graph to orchestrate the previous two, and add some connectors to define how this top-level state is shared between the different graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "95ae7e52-92ed-41a3-88c4-21b6d7c8b041",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
    "\n",
    "supervisor_node = create_team_supervisor(\n",
    "    llm,\n",
    "    \"You are a supervisor tasked with managing a conversation between the\"\n",
    "    \" following teams: {team_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",
    "    [\"Research team\", \"Paper writing team\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4880e573-612f-4d24-97c1-2079382a4a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Top-level graph state\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[List[BaseMessage], operator.add]\n",
    "    next: str\n",
    "\n",
    "\n",
    "def get_last_message(state: State) -> str:\n",
    "    return state[\"messages\"][-1].content\n",
    "\n",
    "\n",
    "def join_graph(response: dict):\n",
    "    return {\"messages\": [response[\"messages\"][-1]]}\n",
    "\n",
    "\n",
    "# Define the graph.\n",
    "super_graph = StateGraph(State)\n",
    "# First add the nodes, which will do the work\n",
    "super_graph.add_node(\"Research team\", get_last_message | research_chain | join_graph)\n",
    "super_graph.add_node(\n",
    "    \"Paper writing team\", get_last_message | authoring_chain | join_graph\n",
    ")\n",
    "super_graph.add_node(\"supervisor\", supervisor_node)\n",
    "\n",
    "# Define the graph connections, which controls how the logic\n",
    "# propagates through the program\n",
    "super_graph.add_edge(\"Research team\", \"supervisor\")\n",
    "super_graph.add_edge(\"Paper writing team\", \"supervisor\")\n",
    "super_graph.add_conditional_edges(\n",
    "    \"supervisor\",\n",
    "    lambda x: x[\"next\"],\n",
    "    {\n",
    "        \"Paper writing team\": \"Paper writing team\",\n",
    "        \"Research team\": \"Research team\",\n",
    "        \"FINISH\": END,\n",
    "    },\n",
    ")\n",
    "super_graph.set_entry_point(\"supervisor\")\n",
    "super_graph = super_graph.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6b8badbf-d728-44bd-a2a7-5b4e587c92fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.', name='Web Scraper')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='Based on the information gathered from the provided URLs, here is a detailed research report on the North American sturgeon along with a chart that summarizes some of the key data:\\n\\n**North American Sturgeon Research Report**\\n\\n**Overview:**\\nNorth American sturgeons are ancient fish that belong to the Acipenseridae family. These bony, bottom-dwelling fish are characterized by their elongated bodies, scutes (bony plates), and barbels near the mouth. They are anadromous, meaning they migrate from saltwater to freshwater to spawn.\\n\\n**Species and Distribution:**\\n- Atlantic Sturgeon (*Acipenser oxyrinchus*): Found along the East Coast from Canada to Florida.\\n- Green Sturgeon (*Acipenser medirostris*): Found on the West Coast from Alaska to California.\\n- Shortnose Sturgeon (*Acipenser brevirostrum*): Distributed along the East Coast in estuarine and riverine habitats.\\n\\n**Conservation Status:**\\n- Atlantic Sturgeon: Multiple distinct population segments (DPSs) listed as endangered or threatened under the Endangered Species Act (ESA).\\n- Green Sturgeon: Southern DPS listed as threatened under the ESA.\\n- Shortnose Sturgeon: Listed as endangered throughout its range under the ESA.\\n\\n**Threats:**\\n- Bycatch in commercial fisheries.\\n- Habitat degradation and loss due to dams, pollution, and development.\\n- Climate change affecting water temperatures and spawning conditions.\\n- Vessel strikes, particularly for Atlantic sturgeon in high-traffic rivers.\\n\\n**Conservation Efforts:**\\nEfforts include habitat restoration, removal of barriers to improve fish passage, protection from bycatch, and public education. Critical habitat has been designated for certain species, and regulations prohibit the retention of green sturgeon in recreational and commercial fisheries.\\n\\n**Chart Information:**\\nThe chart would include species names, conservation status, geographic distribution, and key threats.\\n\\n| Species             | Conservation Status | Distribution                        | Key Threats                                    |\\n|---------------------|---------------------|-------------------------------------|------------------------------------------------|\\n| Atlantic Sturgeon   | Endangered/Threatened | East Coast (Canada to Florida)     | Bycatch, habitat loss, vessel strikes          |\\n| Green Sturgeon      | Southern DPS Threatened | West Coast (Alaska to California) | Bycatch, habitat degradation, climate change   |\\n| Shortnose Sturgeon  | Endangered            | East Coast estuarine and riverine habitats | Bycatch, habitat degradation, dams          |\\n\\n**Conclusion:**\\nNorth American sturgeons face numerous threats that have led to a decline in their populations. However, ongoing conservation efforts, including habitat restoration and protective regulations, are crucial in supporting their recovery and ensuring the survival of these ancient and ecologically important species.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Paper writing team'}}\n",
      "---\n",
      "{'Paper writing team': {'messages': [HumanMessage(content='The contents of the \"AnalysisOfWhispersOfTheWindPoem.txt\" file have been successfully retrieved. It appears to contain the structured paper that was conceptualized earlier, including the introduction, a detailed analysis of the poem\\'s structure and content, exploration of poetic devices and imagery, discussion on the philosophical and environmental implications, and the paper\\'s conclusion.\\n\\nWith this analysis, the paper provides a comprehensive examination of the poem \"Whispers of the Wind,\" capturing the essence of the poem\\'s use of the wind as a symbol to articulate themes related to the natural world. It reflects on how the poet utilizes language and form to convey the power and beauty of nature, as well as the philosophical undertones that underscore human connection with the environment.\\n\\nThe conclusion of the paper emphasizes the significance of the poem in fostering a greater appreciation for nature and highlights the relevance of its message in the context of contemporary environmental discourse. This analysis serves as an important contribution to the understanding of the poem\\'s artistic and thematic dimensions.', name='DocWriter')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'Research team'}}\n",
      "---\n",
      "{'Research team': {'messages': [HumanMessage(content='The contents of the \"AnalysisOfWhispersOfTheWindPoem.txt\" file have been successfully retrieved. It appears to contain the structured paper that was conceptualized earlier, including the introduction, a detailed analysis of the poem\\'s structure and content, exploration of poetic devices and imagery, discussion on the philosophical and environmental implications, and the paper\\'s conclusion.\\n\\nWith this analysis, the paper provides a comprehensive examination of the poem \"Whispers of the Wind,\" capturing the essence of the poem\\'s use of the wind as a symbol to articulate themes related to the natural world. It reflects on how the poet utilizes language and form to convey the power and beauty of nature, as well as the philosophical undertones that underscore human connection with the environment.\\n\\nThe conclusion of the paper emphasizes the significance of the poem in fostering a greater appreciation for nature and highlights the relevance of its message in the context of contemporary environmental discourse. This analysis serves as an important contribution to the understanding of the poem\\'s artistic and thematic dimensions.')]}}\n",
      "---\n",
      "{'supervisor': {'next': 'FINISH'}}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for s in super_graph.stream(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            HumanMessage(\n",
    "                content=\"Write a brief research report on the North American sturgeon. Include a chart.\"\n",
    "            )\n",
    "        ],\n",
    "    },\n",
    "    {\"recursion_limit\": 150},\n",
    "):\n",
    "    if \"__end__\" not in s:\n",
    "        print(s)\n",
    "        print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb1b06f5-648c-459c-84f3-431807373782",
   "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
}