{ "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." ] }, { "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": 3, "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": "code", "execution_count": 4, "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.tools.tavily_search import TavilySearchResults\n", "from langchain_core.tools import tool\n", "\n", "tavily_tool = TavilySearchResults(max_results=5)\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(file_name)\n", " plt.close(fig)\n", " return f'Saved \"{title}\" plot to {file_name}'" ] }, { "cell_type": "code", "execution_count": 5, "id": "6a430af7-8fce-4e66-ba9e-d940c1bc48e8", "metadata": {}, "outputs": [], "source": [ "import operator\n", "from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict\n", "\n", "from langchain.agents import AgentExecutor, create_openai_functions_agent\n", "from langchain_core.messages import BaseMessage, HumanMessage\n", "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n", "from langchain_core.tools import BaseTool\n", "from langchain_experimental.tools import PythonREPLTool\n", "from langchain_openai import ChatOpenAI\n", "\n", "from langgraph.graph import END, StateGraph\n", "\n", "\n", "class AgentState(TypedDict):\n", " messages: Annotated[Sequence[BaseMessage], operator.add]\n", " next: str\n", "\n", "\n", "workflow = StateGraph(AgentState)\n", "\n", "\n", "def create_worker_node(name: str, llm: ChatOpenAI, tools: list, system_prompt: str):\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", " workflow.add_node(name, chain)\n", "\n", "\n", "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n", "\n", "# Note: these worker nodes don't _have_ to be agents. They can be any DAG, tool, or function\n", "create_worker_node(\"Researcher\", llm, [tavily_tool], \"You are a web researcher.\")\n", "create_worker_node(\"Chart Generator\", llm, [create_plot], \"You are a chart generator.\")\n", "# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION\n", "create_worker_node(\n", " \"Coder\",\n", " llm,\n", " [PythonREPLTool()],\n", " \"You may generate safe python code to analyze data.\",\n", ")" ] }, { "cell_type": "markdown", "id": "d6374825-912f-40c9-910d-afa267b401bf", "metadata": {}, "source": [ "Almost done, now we need to create the team supervisor." ] }, { "cell_type": "code", "execution_count": 6, "id": "17c108a0-6dc3-46fd-a5e6-a1fcfad5458a", "metadata": {}, "outputs": [], "source": [ "# So the team supervisor is an LLM node. It just picks the next t\n", "from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser\n", "\n", "\n", "def create_supervisor(members: List[str], llm: ChatOpenAI, system_prompt: str):\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))\n", " if \"members\" in prompt.input_variables:\n", " prompt = prompt.partial(members=\", \".join(members))\n", " chain = (\n", " prompt\n", " | llm.bind_functions(functions=[function_def], function_call=\"route\")\n", " | JsonOutputFunctionsParser()\n", " )\n", " workflow.add_node(\"supervisor\", chain)\n", " conditional_map = {k: k for k in members}\n", " conditional_map[\"FINISH\"] = END\n", "\n", " for member in members:\n", " workflow.add_edge(member, \"supervisor\")\n", " workflow.add_conditional_edges(\"supervisor\", lambda x: x[\"next\"], conditional_map)" ] }, { "cell_type": "code", "execution_count": 7, "id": "14778e86-077b-4e6a-893c-400e59b0cdbf", "metadata": {}, "outputs": [], "source": [ "create_supervisor(\n", " [\"Researcher\", \"Chart Generator\", \"Coder\"],\n", " llm,\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", "\n", "# Finally, add entrypoint\n", "workflow.set_entry_point(\"supervisor\")\n", "\n", "\n", "def enter(text: str) -> dict:\n", " return {\"messages\": [HumanMessage(content=text)]}\n", "\n", "\n", "graph = enter | workflow.compile()" ] }, { "cell_type": "code", "execution_count": null, "id": "56ba78e9-d9c1-457c-a073-d606d5d3e013", "metadata": {}, "outputs": [], "source": [ "results = graph.invoke(\"Code hello world and print it to the terminal\")\n", "results[\"messages\"][-1].pretty_print()" ] }, { "cell_type": "code", "execution_count": 10, "id": "45a92dfd-0e11-47f5-aad4-b68d24990e34", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "================================\u001b[1m Human Message \u001b[0m=================================\n", "\n", "The chart summarizing the data on the 2023 California wildfires has been created successfully. You can see the representation of the statistics mentioned in the summary including the total number of fires, total acres burned, comparison with the five-year average, the size of the largest wildfire, and the number of fatalities.\n", "\n", "![2023 California Wildfires Overview](sandbox:/ca_wildfires_2023_chart.png)\n" ] } ], "source": [ "results = graph.invoke(\n", " \"Write a research summary of CA wildfires in 2023. Include a chart.\"\n", ")\n", "results[\"messages\"][-1].pretty_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.2" } }, "nbformat": 4, "nbformat_minor": 5 }