[Docs] Update notebooks to use START (#902)

This commit is contained in:
William FH
2024-07-01 21:36:34 -07:00
committed by GitHub
parent 267f5e5234
commit 727e63c01e
67 changed files with 1059 additions and 20258 deletions
+11 -184
View File
@@ -26,10 +26,7 @@
"id": "0d30b6f7-3bec-4d9f-af50-43dfdc81ae6c",
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain langchain_openai langchain_experimental langsmith pandas"
]
"source": ["%%capture --no-stderr\n%pip install -U langgraph langchain langchain_openai langchain_experimental langsmith pandas"]
},
{
"cell_type": "code",
@@ -37,24 +34,7 @@
"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.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\""
]
"source": ["import getpass\nimport os\n\n\ndef _set_if_undefined(var: str):\n if not os.environ.get(var):\n os.environ[var] = getpass.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\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""]
},
{
"cell_type": "markdown",
@@ -72,17 +52,7 @@
"id": "f04c6778-403b-4b49-9b93-678e910d5cec",
"metadata": {},
"outputs": [],
"source": [
"from typing import Annotated\n",
"\n",
"from langchain_community.tools.tavily_search import TavilySearchResults\n",
"from langchain_experimental.tools import PythonREPLTool\n",
"\n",
"tavily_tool = TavilySearchResults(max_results=5)\n",
"\n",
"# This executes code locally, which can be unsafe\n",
"python_repl_tool = PythonREPLTool()"
]
"source": ["from typing import Annotated\n\nfrom langchain_community.tools.tavily_search import TavilySearchResults\nfrom langchain_experimental.tools import PythonREPLTool\n\ntavily_tool = TavilySearchResults(max_results=5)\n\n# This executes code locally, which can be unsafe\npython_repl_tool = PythonREPLTool()"]
},
{
"cell_type": "markdown",
@@ -100,28 +70,7 @@
"id": "c4823dd9-26bd-4e1a-8117-b97b2860211a",
"metadata": {},
"outputs": [],
"source": [
"from langchain.agents import AgentExecutor, create_openai_tools_agent\n",
"from langchain_core.messages import BaseMessage, HumanMessage\n",
"from langchain_openai import ChatOpenAI\n",
"\n",
"\n",
"def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):\n",
" # Each worker node will be given a name and some tools.\n",
" prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
" (\n",
" \"system\",\n",
" system_prompt,\n",
" ),\n",
" MessagesPlaceholder(variable_name=\"messages\"),\n",
" MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
" ]\n",
" )\n",
" agent = create_openai_tools_agent(llm, tools, prompt)\n",
" executor = AgentExecutor(agent=agent, tools=tools)\n",
" return executor"
]
"source": ["from langchain.agents import AgentExecutor, create_openai_tools_agent\nfrom langchain_core.messages import BaseMessage, HumanMessage\nfrom langchain_openai import ChatOpenAI\n\n\ndef create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):\n # Each worker node will be given a name and some tools.\n prompt = ChatPromptTemplate.from_messages(\n [\n (\n \"system\",\n system_prompt,\n ),\n MessagesPlaceholder(variable_name=\"messages\"),\n MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n ]\n )\n agent = create_openai_tools_agent(llm, tools, prompt)\n executor = AgentExecutor(agent=agent, tools=tools)\n return executor"]
},
{
"cell_type": "markdown",
@@ -137,11 +86,7 @@
"id": "80862241-a1a7-4726-bce5-f867b233832e",
"metadata": {},
"outputs": [],
"source": [
"def agent_node(state, agent, name):\n",
" result = agent.invoke(state)\n",
" return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}"
]
"source": ["def agent_node(state, agent, name):\n result = agent.invoke(state)\n return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}"]
},
{
"cell_type": "markdown",
@@ -159,59 +104,7 @@
"id": "311f0a58-b425-4496-adac-dc4cd8ffb912",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser\n",
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
"\n",
"members = [\"Researcher\", \"Coder\"]\n",
"system_prompt = (\n",
" \"You are a supervisor tasked with managing a conversation between the\"\n",
" \" following workers: {members}. Given the following user request,\"\n",
" \" respond with the worker to act next. Each worker will perform a\"\n",
" \" task and respond with their results and status. When finished,\"\n",
" \" respond with FINISH.\"\n",
")\n",
"# Our team supervisor is an LLM node. It just picks the next agent to process\n",
"# and decides when the work is completed\n",
"options = [\"FINISH\"] + members\n",
"# Using openai function calling can make output parsing easier for us\n",
"function_def = {\n",
" \"name\": \"route\",\n",
" \"description\": \"Select the next role.\",\n",
" \"parameters\": {\n",
" \"title\": \"routeSchema\",\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"next\": {\n",
" \"title\": \"Next\",\n",
" \"anyOf\": [\n",
" {\"enum\": options},\n",
" ],\n",
" }\n",
" },\n",
" \"required\": [\"next\"],\n",
" },\n",
"}\n",
"prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
" (\"system\", system_prompt),\n",
" MessagesPlaceholder(variable_name=\"messages\"),\n",
" (\n",
" \"system\",\n",
" \"Given the conversation above, who should act next?\"\n",
" \" Or should we FINISH? Select one of: {options}\",\n",
" ),\n",
" ]\n",
").partial(options=str(options), members=\", \".join(members))\n",
"\n",
"llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
"\n",
"supervisor_chain = (\n",
" prompt\n",
" | llm.bind_functions(functions=[function_def], function_call=\"route\")\n",
" | JsonOutputFunctionsParser()\n",
")"
]
"source": ["from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser\nfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n\nmembers = [\"Researcher\", \"Coder\"]\nsystem_prompt = (\n \"You are a supervisor tasked with managing a conversation between the\"\n \" following workers: {members}. Given the following user request,\"\n \" respond with the worker to act next. Each worker will perform a\"\n \" task and respond with their results and status. When finished,\"\n \" respond with FINISH.\"\n)\n# Our team supervisor is an LLM node. It just picks the next agent to process\n# and decides when the work is completed\noptions = [\"FINISH\"] + members\n# Using openai function calling can make output parsing easier for us\nfunction_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}\nprompt = ChatPromptTemplate.from_messages(\n [\n (\"system\", system_prompt),\n MessagesPlaceholder(variable_name=\"messages\"),\n (\n \"system\",\n \"Given the conversation above, who should act next?\"\n \" Or should we FINISH? Select one of: {options}\",\n ),\n ]\n).partial(options=str(options), members=\", \".join(members))\n\nllm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n\nsupervisor_chain = (\n prompt\n | llm.bind_functions(functions=[function_def], function_call=\"route\")\n | JsonOutputFunctionsParser()\n)"]
},
{
"cell_type": "markdown",
@@ -229,41 +122,7 @@
"id": "6a430af7-8fce-4e66-ba9e-d940c1bc48e8",
"metadata": {},
"outputs": [],
"source": [
"import functools\n",
"import operator\n",
"from typing import Sequence, TypedDict\n",
"\n",
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
"\n",
"from langgraph.graph import END, StateGraph\n",
"\n",
"\n",
"# The agent state is the input to each node in the graph\n",
"class AgentState(TypedDict):\n",
" # The annotation tells the graph that new messages will always\n",
" # be added to the current states\n",
" messages: Annotated[Sequence[BaseMessage], operator.add]\n",
" # The 'next' field indicates where to route to next\n",
" next: str\n",
"\n",
"\n",
"research_agent = create_agent(llm, [tavily_tool], \"You are a web researcher.\")\n",
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n",
"\n",
"# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION\n",
"code_agent = create_agent(\n",
" llm,\n",
" [python_repl_tool],\n",
" \"You may generate safe python code to analyze data and generate charts using matplotlib.\",\n",
")\n",
"code_node = functools.partial(agent_node, agent=code_agent, name=\"Coder\")\n",
"\n",
"workflow = StateGraph(AgentState)\n",
"workflow.add_node(\"Researcher\", research_node)\n",
"workflow.add_node(\"Coder\", code_node)\n",
"workflow.add_node(\"supervisor\", supervisor_chain)"
]
"source": ["import functools\nimport operator\nfrom typing import Sequence, TypedDict\n\nfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n\nfrom langgraph.graph import END, StateGraph, START\n\n\n# The agent state is the input to each node in the graph\nclass AgentState(TypedDict):\n # The annotation tells the graph that new messages will always\n # be added to the current states\n messages: Annotated[Sequence[BaseMessage], operator.add]\n # The 'next' field indicates where to route to next\n next: str\n\n\nresearch_agent = create_agent(llm, [tavily_tool], \"You are a web researcher.\")\nresearch_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n\n# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION\ncode_agent = create_agent(\n llm,\n [python_repl_tool],\n \"You may generate safe python code to analyze data and generate charts using matplotlib.\",\n)\ncode_node = functools.partial(agent_node, agent=code_agent, name=\"Coder\")\n\nworkflow = StateGraph(AgentState)\nworkflow.add_node(\"Researcher\", research_node)\nworkflow.add_node(\"Coder\", code_node)\nworkflow.add_node(\"supervisor\", supervisor_chain)"]
},
{
"cell_type": "markdown",
@@ -279,20 +138,7 @@
"id": "14778e86-077b-4e6a-893c-400e59b0cdbf",
"metadata": {},
"outputs": [],
"source": [
"for member in members:\n",
" # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
" workflow.add_edge(member, \"supervisor\")\n",
"# The supervisor populates the \"next\" field in the graph state\n",
"# which routes to a node or finishes\n",
"conditional_map = {k: k for k in members}\n",
"conditional_map[\"FINISH\"] = END\n",
"workflow.add_conditional_edges(\"supervisor\", lambda x: x[\"next\"], conditional_map)\n",
"# Finally, add entrypoint\n",
"workflow.set_entry_point(\"supervisor\")\n",
"\n",
"graph = workflow.compile()"
]
"source": ["for member in members:\n # We want our workers to ALWAYS \"report back\" to the supervisor when done\n workflow.add_edge(member, \"supervisor\")\n# The supervisor populates the \"next\" field in the graph state\n# which routes to a node or finishes\nconditional_map = {k: k for k in members}\nconditional_map[\"FINISH\"] = END\nworkflow.add_conditional_edges(\"supervisor\", lambda x: x[\"next\"], conditional_map)\n# Finally, add entrypoint\nworkflow.add_edge(START, \"supervisor\")\n\ngraph = workflow.compile()"]
},
{
"cell_type": "markdown",
@@ -336,18 +182,7 @@
]
}
],
"source": [
"for s in graph.stream(\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(content=\"Code hello world and print it to the terminal\")\n",
" ]\n",
" }\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"----\")"
]
"source": ["for s in graph.stream(\n {\n \"messages\": [\n HumanMessage(content=\"Code hello world and print it to the terminal\")\n ]\n }\n):\n if \"__end__\" not in s:\n print(s)\n print(\"----\")"]
},
{
"cell_type": "code",
@@ -368,15 +203,7 @@
]
}
],
"source": [
"for s in graph.stream(\n",
" {\"messages\": [HumanMessage(content=\"Write a brief research report on pikas.\")]},\n",
" {\"recursion_limit\": 100},\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"----\")"
]
"source": ["for s in graph.stream(\n {\"messages\": [HumanMessage(content=\"Write a brief research report on pikas.\")]},\n {\"recursion_limit\": 100},\n):\n if \"__end__\" not in s:\n print(s)\n print(\"----\")"]
},
{
"cell_type": "code",
@@ -384,7 +211,7 @@
"id": "1d363d2c-e0da-4cce-ba47-ad2aa9df0fef",
"metadata": {},
"outputs": [],
"source": []
"source": [""]
}
],
"metadata": {
@@ -40,10 +40,7 @@
}
},
"outputs": [],
"source": [
"# %%capture --no-stderr\n",
"# %pip install -U langgraph langchain langchain_openai langchain_experimental"
]
"source": ["# %%capture --no-stderr\n# %pip install -U langgraph langchain langchain_openai langchain_experimental"]
},
{
"cell_type": "code",
@@ -56,25 +53,7 @@
}
},
"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.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\""
]
"source": ["import getpass\nimport os\n\n\ndef _set_if_undefined(var: str):\n if not os.environ.get(var):\n os.environ[var] = getpass.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\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""]
},
{
"cell_type": "markdown",
@@ -103,28 +82,7 @@
}
},
"outputs": [],
"source": [
"from typing import Annotated, List\n",
"\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",
"\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",
" )"
]
"source": ["from typing import Annotated, List\n\nfrom langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.tools.tavily_search import TavilySearchResults\nfrom langchain_core.tools import tool\n\ntavily_tool = TavilySearchResults(max_results=5)\n\n\n@tool\ndef 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",
@@ -150,99 +108,7 @@
}
},
"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}\""
]
"source": ["from pathlib import Path\nfrom tempfile import TemporaryDirectory\nfrom typing import Dict, Optional\n\nfrom langchain_experimental.utilities import PythonREPL\nfrom typing_extensions import TypedDict\n\n_TEMP_DIRECTORY = TemporaryDirectory()\nWORKING_DIRECTORY = Path(_TEMP_DIRECTORY.name)\n\n\n@tool\ndef 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\ndef 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\ndef 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\ndef 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\nrepl = PythonREPL()\n\n\n@tool\ndef 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",
@@ -270,84 +136,7 @@
}
},
"outputs": [],
"source": [
"from typing import List, Optional\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_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",
" )"
]
"source": ["from typing import List, Optional\n\nfrom langchain.agents import AgentExecutor, create_openai_functions_agent\nfrom langchain.output_parsers.openai_functions import JsonOutputFunctionsParser\nfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\nfrom langchain_openai import ChatOpenAI\n\nfrom langgraph.graph import END, StateGraph, START\n\n\ndef 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\ndef agent_node(state, agent, name):\n result = agent.invoke(state)\n return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}\n\n\ndef 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",
@@ -374,52 +163,7 @@
}
},
"outputs": [],
"source": [
"import functools\n",
"import operator\n",
"\n",
"from langchain_core.messages import BaseMessage, HumanMessage\n",
"from langchain_openai.chat_models import ChatOpenAI\n",
"\n",
"\n",
"# ResearchTeam 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=\"WebScraper\")\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, WebScraper. 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\", \"WebScraper\"],\n",
")"
]
"source": ["import functools\nimport operator\n\nfrom langchain_core.messages import BaseMessage, HumanMessage\nfrom langchain_openai.chat_models import ChatOpenAI\n\n\n# ResearchTeam graph state\nclass 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\nllm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n\nsearch_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)\nsearch_node = functools.partial(agent_node, agent=search_agent, name=\"Search\")\n\nresearch_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)\nresearch_node = functools.partial(agent_node, agent=research_agent, name=\"WebScraper\")\n\nsupervisor_agent = create_team_supervisor(\n llm,\n \"You are a supervisor tasked with managing a conversation between the\"\n \" following workers: Search, WebScraper. 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\", \"WebScraper\"],\n)"]
},
{
"cell_type": "markdown",
@@ -440,38 +184,7 @@
}
},
"outputs": [],
"source": [
"research_graph = StateGraph(ResearchTeamState)\n",
"research_graph.add_node(\"Search\", search_node)\n",
"research_graph.add_node(\"WebScraper\", 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(\"WebScraper\", \"supervisor\")\n",
"research_graph.add_conditional_edges(\n",
" \"supervisor\",\n",
" lambda x: x[\"next\"],\n",
" {\"Search\": \"Search\", \"WebScraper\": \"WebScraper\", \"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"
]
"source": ["research_graph = StateGraph(ResearchTeamState)\nresearch_graph.add_node(\"Search\", search_node)\nresearch_graph.add_node(\"WebScraper\", research_node)\nresearch_graph.add_node(\"supervisor\", supervisor_agent)\n\n# Define the control flow\nresearch_graph.add_edge(\"Search\", \"supervisor\")\nresearch_graph.add_edge(\"WebScraper\", \"supervisor\")\nresearch_graph.add_conditional_edges(\n \"supervisor\",\n lambda x: x[\"next\"],\n {\"Search\": \"Search\", \"WebScraper\": \"WebScraper\", \"FINISH\": END},\n)\n\n\nresearch_graph.add_edge(START, \"supervisor\")\nchain = 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\ndef enter_chain(message: str):\n results = {\n \"messages\": [HumanMessage(content=message)],\n }\n return results\n\n\nresearch_chain = enter_chain | chain"]
},
{
"cell_type": "code",
@@ -495,11 +208,7 @@
"output_type": "display_data"
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"display(Image(chain.get_graph(xray=True).draw_mermaid_png()))"
]
"source": ["from IPython.display import Image, display\n\ndisplay(Image(chain.get_graph(xray=True).draw_mermaid_png()))"]
},
{
"cell_type": "markdown",
@@ -520,14 +229,7 @@
}
},
"outputs": [],
"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(\"---\")"
]
"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",
@@ -552,92 +254,7 @@
}
},
"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 Exception:\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",
")"
]
"source": ["import operator\nfrom pathlib import Path\n\n\n# Document writing team graph state\nclass 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.\ndef 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 Exception:\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\nllm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n\ndoc_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\ncontext_aware_doc_writer_agent = prelude | doc_writer_agent\ndoc_writing_node = functools.partial(\n agent_node, agent=context_aware_doc_writer_agent, name=\"DocWriter\"\n)\n\nnote_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)\ncontext_aware_note_taking_agent = prelude | note_taking_agent\nnote_taking_node = functools.partial(\n agent_node, agent=context_aware_note_taking_agent, name=\"NoteTaker\"\n)\n\nchart_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)\ncontext_aware_chart_generating_agent = prelude | chart_generating_agent\nchart_generating_node = functools.partial(\n agent_node, agent=context_aware_note_taking_agent, name=\"ChartGenerator\"\n)\n\ndoc_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",
@@ -658,53 +275,7 @@
}
},
"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",
")"
]
"source": ["# Create the graph here:\n# Note that we have unrolled the loop for the sake of this doc\nauthoring_graph = StateGraph(DocWritingState)\nauthoring_graph.add_node(\"DocWriter\", doc_writing_node)\nauthoring_graph.add_node(\"NoteTaker\", note_taking_node)\nauthoring_graph.add_node(\"ChartGenerator\", chart_generating_node)\nauthoring_graph.add_node(\"supervisor\", doc_writing_supervisor)\n\n# Add the edges that always occur\nauthoring_graph.add_edge(\"DocWriter\", \"supervisor\")\nauthoring_graph.add_edge(\"NoteTaker\", \"supervisor\")\nauthoring_graph.add_edge(\"ChartGenerator\", \"supervisor\")\n\n# Add the edges where routing applies\nauthoring_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\nauthoring_graph.add_edge(START, \"supervisor\")\nchain = 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\ndef 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\nauthoring_chain = (\n functools.partial(enter_chain, members=authoring_graph.nodes)\n | authoring_graph.compile()\n)"]
},
{
"cell_type": "code",
@@ -728,11 +299,7 @@
"output_type": "display_data"
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"display(Image(chain.get_graph().draw_mermaid_png()))"
]
"source": ["from IPython.display import Image, display\n\ndisplay(Image(chain.get_graph().draw_mermaid_png()))"]
},
{
"cell_type": "code",
@@ -745,15 +312,7 @@
}
},
"outputs": [],
"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(\"---\")"
]
"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",
@@ -778,22 +337,7 @@
}
},
"outputs": [],
"source": [
"from langchain_core.messages import BaseMessage\n",
"from langchain_openai.chat_models import ChatOpenAI\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",
" [\"ResearchTeam\", \"PaperWritingTeam\"],\n",
")"
]
"source": ["from langchain_core.messages import BaseMessage\nfrom langchain_openai.chat_models import ChatOpenAI\n\nllm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n\nsupervisor_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 [\"ResearchTeam\", \"PaperWritingTeam\"],\n)"]
},
{
"cell_type": "code",
@@ -806,46 +350,7 @@
}
},
"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(\"ResearchTeam\", get_last_message | research_chain | join_graph)\n",
"super_graph.add_node(\n",
" \"PaperWritingTeam\", 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(\"ResearchTeam\", \"supervisor\")\n",
"super_graph.add_edge(\"PaperWritingTeam\", \"supervisor\")\n",
"super_graph.add_conditional_edges(\n",
" \"supervisor\",\n",
" lambda x: x[\"next\"],\n",
" {\n",
" \"PaperWritingTeam\": \"PaperWritingTeam\",\n",
" \"ResearchTeam\": \"ResearchTeam\",\n",
" \"FINISH\": END,\n",
" },\n",
")\n",
"super_graph.set_entry_point(\"supervisor\")\n",
"super_graph = super_graph.compile()"
]
"source": ["# Top-level graph state\nclass State(TypedDict):\n messages: Annotated[List[BaseMessage], operator.add]\n next: str\n\n\ndef get_last_message(state: State) -> str:\n return state[\"messages\"][-1].content\n\n\ndef join_graph(response: dict):\n return {\"messages\": [response[\"messages\"][-1]]}\n\n\n# Define the graph.\nsuper_graph = StateGraph(State)\n# First add the nodes, which will do the work\nsuper_graph.add_node(\"ResearchTeam\", get_last_message | research_chain | join_graph)\nsuper_graph.add_node(\n \"PaperWritingTeam\", get_last_message | authoring_chain | join_graph\n)\nsuper_graph.add_node(\"supervisor\", supervisor_node)\n\n# Define the graph connections, which controls how the logic\n# propagates through the program\nsuper_graph.add_edge(\"ResearchTeam\", \"supervisor\")\nsuper_graph.add_edge(\"PaperWritingTeam\", \"supervisor\")\nsuper_graph.add_conditional_edges(\n \"supervisor\",\n lambda x: x[\"next\"],\n {\n \"PaperWritingTeam\": \"PaperWritingTeam\",\n \"ResearchTeam\": \"ResearchTeam\",\n \"FINISH\": END,\n },\n)\nsuper_graph.add_edge(START, \"supervisor\")\nsuper_graph = super_graph.compile()"]
},
{
"cell_type": "code",
@@ -869,11 +374,7 @@
"output_type": "display_data"
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"display(Image(super_graph.get_graph().draw_mermaid_png()))"
]
"source": ["from IPython.display import Image, display\n\ndisplay(Image(super_graph.get_graph().draw_mermaid_png()))"]
},
{
"cell_type": "code",
@@ -886,21 +387,7 @@
}
},
"outputs": [],
"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(\"---\")"
]
"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(\"---\")"]
}
],
"metadata": {
@@ -26,10 +26,7 @@
"id": "0d7b6dcc-c985-46e2-8457-7e6b0298b950",
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langchain langchain_openai langsmith pandas langchain_experimental matplotlib langgraph langchain_core"
]
"source": ["%%capture --no-stderr\n%pip install -U langchain langchain_openai langsmith pandas langchain_experimental matplotlib langgraph langchain_core"]
},
{
"cell_type": "code",
@@ -37,24 +34,7 @@
"id": "743c19df-6da9-4d1e-b2d2-ea40080b9fdc",
"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.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\""
]
"source": ["import getpass\nimport os\n\n\ndef _set_if_undefined(var: str):\n if not os.environ.get(var):\n os.environ[var] = getpass.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\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""]
},
{
"cell_type": "markdown",
@@ -74,38 +54,7 @@
"id": "4325a10e-38dc-4a98-9004-e1525eaba377",
"metadata": {},
"outputs": [],
"source": [
"from langchain_core.messages import (\n",
" BaseMessage,\n",
" HumanMessage,\n",
" ToolMessage,\n",
")\n",
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
"\n",
"from langgraph.graph import END, StateGraph\n",
"\n",
"\n",
"def create_agent(llm, tools, system_message: str):\n",
" \"\"\"Create an agent.\"\"\"\n",
" prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
" (\n",
" \"system\",\n",
" \"You are a helpful AI assistant, collaborating with other assistants.\"\n",
" \" Use the provided tools to progress towards answering the question.\"\n",
" \" If you are unable to fully answer, that's OK, another assistant with different tools \"\n",
" \" will help where you left off. Execute what you can to make progress.\"\n",
" \" If you or any of the other assistants have the final answer or deliverable,\"\n",
" \" prefix your response with FINAL ANSWER so the team knows to stop.\"\n",
" \" You have access to the following tools: {tool_names}.\\n{system_message}\",\n",
" ),\n",
" MessagesPlaceholder(variable_name=\"messages\"),\n",
" ]\n",
" )\n",
" prompt = prompt.partial(system_message=system_message)\n",
" prompt = prompt.partial(tool_names=\", \".join([tool.name for tool in tools]))\n",
" return prompt | llm.bind_tools(tools)"
]
"source": ["from langchain_core.messages import (\n BaseMessage,\n HumanMessage,\n ToolMessage,\n)\nfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n\nfrom langgraph.graph import END, StateGraph, START\n\n\ndef create_agent(llm, tools, system_message: str):\n \"\"\"Create an agent.\"\"\"\n prompt = ChatPromptTemplate.from_messages(\n [\n (\n \"system\",\n \"You are a helpful AI assistant, collaborating with other assistants.\"\n \" Use the provided tools to progress towards answering the question.\"\n \" If you are unable to fully answer, that's OK, another assistant with different tools \"\n \" will help where you left off. Execute what you can to make progress.\"\n \" If you or any of the other assistants have the final answer or deliverable,\"\n \" prefix your response with FINAL ANSWER so the team knows to stop.\"\n \" You have access to the following tools: {tool_names}.\\n{system_message}\",\n ),\n MessagesPlaceholder(variable_name=\"messages\"),\n ]\n )\n prompt = prompt.partial(system_message=system_message)\n prompt = prompt.partial(tool_names=\", \".join([tool.name for tool in tools]))\n return prompt | llm.bind_tools(tools)"]
},
{
"cell_type": "markdown",
@@ -123,35 +72,7 @@
"id": "ca076f3b-a729-4ca9-8f91-05c2ba58d610",
"metadata": {},
"outputs": [],
"source": [
"from typing import Annotated\n",
"\n",
"from langchain_community.tools.tavily_search import TavilySearchResults\n",
"from langchain_core.tools import tool\n",
"from langchain_experimental.utilities import PythonREPL\n",
"\n",
"tavily_tool = TavilySearchResults(max_results=5)\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",
" result_str = f\"Successfully executed:\\n```python\\n{code}\\n```\\nStdout: {result}\"\n",
" return (\n",
" result_str + \"\\n\\nIf you have completed all tasks, respond with FINAL ANSWER.\"\n",
" )"
]
"source": ["from typing import Annotated\n\nfrom langchain_community.tools.tavily_search import TavilySearchResults\nfrom langchain_core.tools import tool\nfrom langchain_experimental.utilities import PythonREPL\n\ntavily_tool = TavilySearchResults(max_results=5)\n\n# Warning: This executes code locally, which can be unsafe when not sandboxed\n\nrepl = PythonREPL()\n\n\n@tool\ndef 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 result_str = f\"Successfully executed:\\n```python\\n{code}\\n```\\nStdout: {result}\"\n return (\n result_str + \"\\n\\nIf you have completed all tasks, respond with FINAL ANSWER.\"\n )"]
},
{
"cell_type": "markdown",
@@ -179,19 +100,7 @@
"id": "290c91d4-f6f4-443c-8181-233d39102974",
"metadata": {},
"outputs": [],
"source": [
"import operator\n",
"from typing import Annotated, Sequence, TypedDict\n",
"\n",
"from langchain_openai import ChatOpenAI\n",
"\n",
"\n",
"# This defines the object that is passed between each node\n",
"# in the graph. We will create different nodes for each agent and tool\n",
"class AgentState(TypedDict):\n",
" messages: Annotated[Sequence[BaseMessage], operator.add]\n",
" sender: str"
]
"source": ["import operator\nfrom typing import Annotated, Sequence, TypedDict\n\nfrom langchain_openai import ChatOpenAI\n\n\n# This defines the object that is passed between each node\n# in the graph. We will create different nodes for each agent and tool\nclass AgentState(TypedDict):\n messages: Annotated[Sequence[BaseMessage], operator.add]\n sender: str"]
},
{
"cell_type": "markdown",
@@ -209,46 +118,7 @@
"id": "71b790ca-9cef-4b22-b469-4b1d5d8424d6",
"metadata": {},
"outputs": [],
"source": [
"import functools\n",
"\n",
"from langchain_core.messages import AIMessage\n",
"\n",
"\n",
"# Helper function to create a node for a given agent\n",
"def agent_node(state, agent, name):\n",
" result = agent.invoke(state)\n",
" # We convert the agent output into a format that is suitable to append to the global state\n",
" if isinstance(result, ToolMessage):\n",
" pass\n",
" else:\n",
" result = AIMessage(**result.dict(exclude={\"type\", \"name\"}), name=name)\n",
" return {\n",
" \"messages\": [result],\n",
" # Since we have a strict workflow, we can\n",
" # track the sender so we know who to pass to next.\n",
" \"sender\": name,\n",
" }\n",
"\n",
"\n",
"llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
"\n",
"# Research agent and node\n",
"research_agent = create_agent(\n",
" llm,\n",
" [tavily_tool],\n",
" system_message=\"You should provide accurate data for the chart_generator to use.\",\n",
")\n",
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n",
"\n",
"# chart_generator\n",
"chart_agent = create_agent(\n",
" llm,\n",
" [python_repl],\n",
" system_message=\"Any charts you display will be visible by the user.\",\n",
")\n",
"chart_node = functools.partial(agent_node, agent=chart_agent, name=\"chart_generator\")"
]
"source": ["import functools\n\nfrom langchain_core.messages import AIMessage\n\n\n# Helper function to create a node for a given agent\ndef agent_node(state, agent, name):\n result = agent.invoke(state)\n # We convert the agent output into a format that is suitable to append to the global state\n if isinstance(result, ToolMessage):\n pass\n else:\n result = AIMessage(**result.dict(exclude={\"type\", \"name\"}), name=name)\n return {\n \"messages\": [result],\n # Since we have a strict workflow, we can\n # track the sender so we know who to pass to next.\n \"sender\": name,\n }\n\n\nllm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n\n# Research agent and node\nresearch_agent = create_agent(\n llm,\n [tavily_tool],\n system_message=\"You should provide accurate data for the chart_generator to use.\",\n)\nresearch_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n\n# chart_generator\nchart_agent = create_agent(\n llm,\n [python_repl],\n system_message=\"Any charts you display will be visible by the user.\",\n)\nchart_node = functools.partial(agent_node, agent=chart_agent, name=\"chart_generator\")"]
},
{
"cell_type": "markdown",
@@ -266,12 +136,7 @@
"id": "d9a79c76-5c7c-42f6-91cf-635bc8305804",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.prebuilt import ToolNode\n",
"\n",
"tools = [tavily_tool, python_repl]\n",
"tool_node = ToolNode(tools)"
]
"source": ["from langgraph.prebuilt import ToolNode\n\ntools = [tavily_tool, python_repl]\ntool_node = ToolNode(tools)"]
},
{
"cell_type": "markdown",
@@ -289,23 +154,7 @@
"id": "4f4b4d37-e8a3-4abb-8d42-eaea26016f35",
"metadata": {},
"outputs": [],
"source": [
"# Either agent can decide to end\n",
"from typing import Literal\n",
"\n",
"\n",
"def router(state) -> Literal[\"call_tool\", \"__end__\", \"continue\"]:\n",
" # This is the router\n",
" messages = state[\"messages\"]\n",
" last_message = messages[-1]\n",
" if last_message.tool_calls:\n",
" # The previous agent is invoking a tool\n",
" return \"call_tool\"\n",
" if \"FINAL ANSWER\" in last_message.content:\n",
" # Any agent decided the work is done\n",
" return \"__end__\"\n",
" return \"continue\""
]
"source": ["# Either agent can decide to end\nfrom typing import Literal\n\n\ndef router(state) -> Literal[\"call_tool\", \"__end__\", \"continue\"]:\n # This is the router\n messages = state[\"messages\"]\n last_message = messages[-1]\n if last_message.tool_calls:\n # The previous agent is invoking a tool\n return \"call_tool\"\n if \"FINAL ANSWER\" in last_message.content:\n # Any agent decided the work is done\n return \"__end__\"\n return \"continue\""]
},
{
"cell_type": "markdown",
@@ -323,39 +172,7 @@
"id": "4dce3901-6ad5-4df5-8528-6e865cf96cb0",
"metadata": {},
"outputs": [],
"source": [
"workflow = StateGraph(AgentState)\n",
"\n",
"workflow.add_node(\"Researcher\", research_node)\n",
"workflow.add_node(\"chart_generator\", chart_node)\n",
"workflow.add_node(\"call_tool\", tool_node)\n",
"\n",
"workflow.add_conditional_edges(\n",
" \"Researcher\",\n",
" router,\n",
" {\"continue\": \"chart_generator\", \"call_tool\": \"call_tool\", \"__end__\": END},\n",
")\n",
"workflow.add_conditional_edges(\n",
" \"chart_generator\",\n",
" router,\n",
" {\"continue\": \"Researcher\", \"call_tool\": \"call_tool\", \"__end__\": END},\n",
")\n",
"\n",
"workflow.add_conditional_edges(\n",
" \"call_tool\",\n",
" # Each agent node updates the 'sender' field\n",
" # the tool calling node does not, meaning\n",
" # this edge will route back to the original agent\n",
" # who invoked the tool\n",
" lambda x: x[\"sender\"],\n",
" {\n",
" \"Researcher\": \"Researcher\",\n",
" \"chart_generator\": \"chart_generator\",\n",
" },\n",
")\n",
"workflow.set_entry_point(\"Researcher\")\n",
"graph = workflow.compile()"
]
"source": ["workflow = StateGraph(AgentState)\n\nworkflow.add_node(\"Researcher\", research_node)\nworkflow.add_node(\"chart_generator\", chart_node)\nworkflow.add_node(\"call_tool\", tool_node)\n\nworkflow.add_conditional_edges(\n \"Researcher\",\n router,\n {\"continue\": \"chart_generator\", \"call_tool\": \"call_tool\", \"__end__\": END},\n)\nworkflow.add_conditional_edges(\n \"chart_generator\",\n router,\n {\"continue\": \"Researcher\", \"call_tool\": \"call_tool\", \"__end__\": END},\n)\n\nworkflow.add_conditional_edges(\n \"call_tool\",\n # Each agent node updates the 'sender' field\n # the tool calling node does not, meaning\n # this edge will route back to the original agent\n # who invoked the tool\n lambda x: x[\"sender\"],\n {\n \"Researcher\": \"Researcher\",\n \"chart_generator\": \"chart_generator\",\n },\n)\nworkflow.add_edge(START, \"Researcher\")\ngraph = workflow.compile()"]
},
{
"cell_type": "code",
@@ -374,15 +191,7 @@
"output_type": "display_data"
}
],
"source": [
"from IPython.display import Image, display\n",
"\n",
"try:\n",
" display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
"except Exception:\n",
" # This requires some extra dependencies and is optional\n",
" pass"
]
"source": ["from IPython.display import Image, display\n\ntry:\n display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\nexcept Exception:\n # This requires some extra dependencies and is optional\n pass"]
},
{
"cell_type": "markdown",
@@ -481,24 +290,7 @@
]
}
],
"source": [
"events = graph.stream(\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(\n",
" content=\"Fetch the UK's GDP over the past 5 years,\"\n",
" \" then draw a line graph of it.\"\n",
" \" Once you code it up, finish.\"\n",
" )\n",
" ],\n",
" },\n",
" # Maximum number of steps to take in the graph\n",
" {\"recursion_limit\": 150},\n",
")\n",
"for s in events:\n",
" print(s)\n",
" print(\"----\")"
]
"source": ["events = graph.stream(\n {\n \"messages\": [\n HumanMessage(\n content=\"Fetch the UK's GDP over the past 5 years,\"\n \" then draw a line graph of it.\"\n \" Once you code it up, finish.\"\n )\n ],\n },\n # Maximum number of steps to take in the graph\n {\"recursion_limit\": 150},\n)\nfor s in events:\n print(s)\n print(\"----\")"]
},
{
"cell_type": "code",
@@ -506,7 +298,7 @@
"id": "010fc36e-4116-4758-bcac-b02c7dcd405d",
"metadata": {},
"outputs": [],
"source": []
"source": [""]
}
],
"metadata": {