diff --git a/examples/advanced_agents/multi-agent/agent_supervisor.ipynb b/examples/advanced_agents/multi-agent/agent_supervisor.ipynb index 1723198b4..21acbcbc5 100644 --- a/examples/advanced_agents/multi-agent/agent_supervisor.ipynb +++ b/examples/advanced_agents/multi-agent/agent_supervisor.ipynb @@ -5,9 +5,9 @@ "id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276", "metadata": {}, "source": [ - "## Multi-agent Example 2: Agent Team Supervisor\n", + "## Multi-agent Example 3: Agent Team Supervisor\n", "\n", - "The prevoius example routed messages automatically based on the output of the initial researcher agent.\n", + "The previous example 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", diff --git a/examples/advanced_agents/multi-agent/hierarchical_agent_teams.ipynb b/examples/advanced_agents/multi-agent/hierarchical_agent_teams.ipynb new file mode 100644 index 000000000..64942f416 --- /dev/null +++ b/examples/advanced_agents/multi-agent/hierarchical_agent_teams.ipynb @@ -0,0 +1,637 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276", + "metadata": {}, + "source": [ + "## 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? Or 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", + "\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", + "\n", + "\n", + "In the rest of this notebook, you will:\n", + "1. Define some utilities to help create the graph and their relations\n", + "2. Write the tools and agent implementations.\n", + "3. Generate each team's sub-graph.\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 langchain langchain_openai langsmith pandas" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "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", + "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n", + "os.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\"" + ] + }, + { + "cell_type": "markdown", + "id": "504ee1c6-2b6a-439d-9046-df54e1e15698", + "metadata": {}, + "source": [ + "## Define Utilities\n", + "\n", + "We are going to create a few utility functions to help us:\n", + "\n", + "1. Create an agent and add it to a graph.\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": 7, + "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", + "class WorkerAgent(TypedDict):\n", + " name: str\n", + " description: str\n", + "\n", + "\n", + "def create_worker_agent(\n", + " graph_builder: StateGraph,\n", + " name: str,\n", + " llm: ChatOpenAI,\n", + " tools: list,\n", + " system_prompt: str,\n", + " prelude: Optional[Union[Runnable, Callable]] = None, # Optional required steps\n", + ") -> str:\n", + " \"\"\"Create a function-calling agent and add it to the graph.\"\"\"\n", + " system_prompt += \"\\nYou 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", + " chain = executor | (\n", + " lambda x: {\"messages\": [HumanMessage(content=x[\"output\"], name=name)]}\n", + " )\n", + " if prelude is not None:\n", + " chain = prelude | chain\n", + " graph_builder.add_node(name, chain)\n", + " return name\n", + "\n", + "\n", + "def create_team_supervisor(\n", + " graph_builder: StateGraph, llm: ChatOpenAI, system_prompt: str\n", + ") -> str:\n", + " \"\"\"An LLM-based router.\"\"\"\n", + " supervisor_id = uuid.uuid4().hex[:4]\n", + " supervisor_name = f\"supervisor - {supervisor_id}\"\n", + " members = list(graph_builder.nodes)\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", + " chain = (\n", + " prompt\n", + " | llm.bind_functions(functions=[function_def], function_call=\"route\")\n", + " | JsonOutputFunctionsParser()\n", + " )\n", + " graph_builder.add_node(supervisor_name, chain)\n", + " conditional_map = {k: k for k in members}\n", + " conditional_map[\"FINISH\"] = END\n", + "\n", + " for member in members:\n", + " graph_builder.add_edge(member, supervisor_name)\n", + " graph_builder.add_conditional_edges(\n", + " supervisor_name, lambda x: x[\"next\"], conditional_map\n", + " )\n", + " return supervisor_name" + ] + }, + { + "cell_type": "markdown", + "id": "00282b1f-bb4d-4ee7-9bae-e8e6f586f12e", + "metadata": {}, + "source": [ + "## Define agents + tools\n", + "\n", + "### Research Team\n", + "\n", + "The research team can use a search engine and web scraper to find information on the web." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f04c6778-403b-4b49-9b93-678e910d5cec", + "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'\\n{doc.page_content}\\n'\n", + " for doc in docs\n", + " ]\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "53db0c78-e357-48ba-ae5f-3fc04735a3b7", + "metadata": {}, + "outputs": [], + "source": [ + "import operator\n", + "\n", + "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage\n", + "from langchain_openai.chat_models import ChatOpenAI\n", + "\n", + "\n", + "# Research team graph\n", + "class State(TypedDict):\n", + " messages: Annotated[List[BaseMessage], operator.add]\n", + " team_members: List[str]\n", + " final_response: AIMessage\n", + " next: str\n", + "\n", + "\n", + "research_graph = StateGraph(State)\n", + "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n", + "create_worker_agent(\n", + " research_graph,\n", + " \"Search\",\n", + " llm,\n", + " [tavily_tool],\n", + " \"You are a research assistant who can search for things using a search engine.\",\n", + ")\n", + "create_worker_agent(\n", + " research_graph,\n", + " \"Web Scraper\",\n", + " llm,\n", + " [tavily_tool],\n", + " \"You are a research assistant who can scrape specified urls for more detailed information.\",\n", + ")\n", + "supervisor_node = create_team_supervisor(\n", + " research_graph,\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", + ")\n", + "\n", + "research_graph.set_entry_point(supervisor_node)\n", + "\n", + "\n", + "def return_final_response(state):\n", + " return {\"final_response\": state[\"messages\"][-1]}\n", + "\n", + "\n", + "research_chain = research_graph.compile() | return_final_response" + ] + }, + { + "cell_type": "markdown", + "id": "749b99ab-f6f0-4c5d-a90b-10102465d186", + "metadata": {}, + "source": [ + "## Document Writing Team\n", + "\n", + "We will construct a graph in a similar fashion. This time using different 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": 10, + "id": "202806d6-80bf-4153-ac16-ed6059236f2a", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from tempfile import TemporaryDirectory\n", + "from typing import Dict\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", + " subpoints: Annotated[\n", + " List[List[str]],\n", + " \"List of lists, each containing subpoints for the corresponding main point.\",\n", + " ],\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", + " if len(points) != len(subpoints):\n", + " raise ValueError(\"Each main point must have a corresponding list of subpoints.\")\n", + "\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", + " for j, subpoint in enumerate(subpoints[i]):\n", + " file.write(f\"\\t{j + 1}. {subpoint}\\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", + " # Read the contents of the file\n", + "\n", + " with (WORKING_DIRECTORY / file_name).open(\"r\") as file:\n", + " lines = file.readlines()\n", + "\n", + " # Adjust the line numbers for 0-indexing and sort\n", + " sorted_inserts = sorted(inserts.items())\n", + "\n", + " # Perform the insertions\n", + " for line_number, text in sorted_inserts:\n", + " if 1 <= line_number <= len(lines) + 1:\n", + " # Insert the text at the specified line number\n", + " lines.insert(line_number - 1, text + \"\\n\")\n", + " else:\n", + " return f\"Error: Line number {line_number} is out of range.\"\n", + "\n", + " # Write the modified content back to the file\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", + "@tool\n", + "def create_plot(\n", + " data: Annotated[\n", + " Union[List[float], List[int]],\n", + " \"Numerical values for bar heights or line points.\",\n", + " ],\n", + " file_name: Annotated[str, \"File path to save the figure.\"],\n", + " labels: Annotated[\n", + " Union[List[str], None], \"Bar or point labels, defaults to None.\"\n", + " ] = None,\n", + " title: Annotated[str, \"Title of the plot.\"] = \"Plot\",\n", + " xlabel: Annotated[str, \"Label for the X-axis.\"] = \"X\",\n", + " ylabel: Annotated[str, \"Label for the Y-axis.\"] = \"Y\",\n", + " color: Annotated[Union[str, List[str]], \"Color(s) for the bars or line.\"] = \"blue\",\n", + " plot_type: Annotated[str, \"Type of plot ('bar' or 'line').\"] = \"bar\",\n", + ") -> Annotated[str, \"Path of the saved figure file.\"]:\n", + " \"\"\"Create a line or bar chart.\"\"\"\n", + " if plot_type not in [\"bar\", \"line\"]:\n", + " raise ValueError(\"Invalid plot_type. Expected 'bar' or 'line'.\")\n", + "\n", + " fig, ax = plt.subplots(figsize=(10, 6))\n", + " x_positions = range(len(data))\n", + "\n", + " if labels and len(labels) == len(data):\n", + " plt.xticks(x_positions, labels)\n", + "\n", + " if plot_type == \"bar\":\n", + " ax.bar(x_positions, data, color=color)\n", + " elif plot_type == \"line\":\n", + " ax.plot(x_positions, data, color=color, marker=\"o\") # 'o' for circular markers\n", + "\n", + " ax.set_title(title)\n", + " ax.set_xlabel(xlabel)\n", + " ax.set_ylabel(ylabel)\n", + " fig.savefig(str(WORKING_DIRECTORY / file_name))\n", + " plt.close(fig)\n", + " return f'Saved \"{title}\" plot to {file_name}'" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "1bcdbf44-9481-430c-8429-fa142ed8a626", + "metadata": {}, + "outputs": [], + "source": [ + "import operator\n", + "from pathlib import Path\n", + "\n", + "\n", + "# Research team graph\n", + "class AuthoringState(TypedDict):\n", + " messages: Annotated[List[BaseMessage], operator.add]\n", + " team_members: str\n", + " final_response: AIMessage\n", + " next: str\n", + " current_files: str\n", + "\n", + "\n", + "authoring_graph = StateGraph(AuthoringState)\n", + "\n", + "\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", + "create_worker_agent(\n", + " authoring_graph,\n", + " \"Author Docs\",\n", + " llm,\n", + " [write_document, edit_document, read_document],\n", + " \"You are an expert writing a research document.\\n\"\n", + " \"Below are files currently in your directory:\\n{current_files}\",\n", + " prelude=prelude,\n", + ")\n", + "create_worker_agent(\n", + " authoring_graph,\n", + " \"Outline + Notetaker\",\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", + " prelude=prelude,\n", + ")\n", + "create_worker_agent(\n", + " authoring_graph,\n", + " \"Generate Charts\",\n", + " llm,\n", + " [read_document, create_plot],\n", + " \"You are a data viz expert tasked with generating charts for a research project.\"\n", + " \"{current_files}\",\n", + ")\n", + "\n", + "supervisor_node = create_team_supervisor(\n", + " authoring_graph,\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", + ")\n", + "\n", + "authoring_graph.set_entry_point(supervisor_node)\n", + "\n", + "\n", + "def enter_chain(message: str):\n", + " return {\n", + " \"messages\": [HumanMessage(content=\"Write a short report on pikas\")],\n", + " \"team_members\": \"\\n\".join(sorted(authoring_graph.nodes)),\n", + " }\n", + "\n", + "\n", + "def return_final_response(state):\n", + " return {\"final_response\": state[\"messages\"][-1]}\n", + "\n", + "\n", + "authoring_chain = enter_chain | authoring_graph.compile() | return_final_response" + ] + }, + { + "cell_type": "markdown", + "id": "f4b5b08d-9a9a-474a-94b4-f7aaa8ff19e6", + "metadata": {}, + "source": [ + "## Adding Hierarchy\n", + "\n", + "We've created two graphs already. Now let's put them together.\n", + "\n", + "We'll create a third graph to orchestrate the previous two, and add some connectors to define how state is shared between the different graphs." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "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", + "# Research team graph\n", + "class State(TypedDict):\n", + " messages: Annotated[List[BaseMessage], operator.add]\n", + " next: str\n", + "\n", + "\n", + "def join_graph(response):\n", + " return {\"messages\": [response[\"final_response\"]]}\n", + "\n", + "\n", + "super_graph = StateGraph(State)\n", + "super_graph.add_node(\"Research team\", research_chain | join_graph)\n", + "super_graph.add_node(\"Paper writing team\", authoring_chain | join_graph)\n", + "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n", + "supervisor_node = create_team_supervisor(\n", + " super_graph,\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", + ")\n", + "\n", + "super_graph.set_entry_point(supervisor_node)\n", + "super_graph = super_graph.compile()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6b8badbf-d728-44bd-a2a7-5b4e587c92fe", + "metadata": {}, + "outputs": [], + "source": [ + "results = super_graph.invoke(\n", + " {\n", + " \"messages\": [\n", + " HumanMessage(\n", + " content=\"Research and write a report about the climate impacts\"\n", + " \" on crop yields in Bangladesh in 2023. Write the paper and include plots.\",\n", + " )\n", + " ]\n", + " },\n", + " # {\"recursion_limit\": 150},\n", + ")\n", + "results[\"messages\"][-1]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6ffcc7f-7b78-4ca5-8e0a-7c0ac08300fc", + "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.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/advanced_agents/multi-agent/img/hierarchical-diagram.png b/examples/advanced_agents/multi-agent/img/hierarchical-diagram.png new file mode 100644 index 000000000..b3fa81275 Binary files /dev/null and b/examples/advanced_agents/multi-agent/img/hierarchical-diagram.png differ