mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Add hierarchical example
This commit is contained in:
William Fu-Hinthorn
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
"\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'<Document name=\"{doc.metadata[\"title\"]}\">\\n{doc.page_content}\\n</Document>'\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
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 108 KiB |
Reference in New Issue
Block a user