diff --git a/examples/advanced_agents/LLMCompiler.ipynb b/examples/advanced_agents/LLMCompiler.ipynb deleted file mode 100644 index 35acbe368..000000000 --- a/examples/advanced_agents/LLMCompiler.ipynb +++ /dev/null @@ -1,1090 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "0c8b472b-f3fb-46c2-841f-930a4692697b", - "metadata": {}, - "source": [ - "# Implementing LLMCompiler using LangGraph\n", - "By Kim, et. al [🔗](https://arxiv.org/abs/2312.04511)\n", - "\n", - "LLMCompiler is an agent architecture intented on speeding up the latency of agentic tasks via fast, parallel tool execution. It has 3 main components:\n", - "\n", - "1. Planner: generate a DAG of tasks.\n", - "2. Task Fetching Unit: schedules and executes the tasks\n", - "3. Joiner: Responds to the user or triggers a second plan\n", - "\n", - "\n", - "This notebook walks through each component and shows how to wire them together using LangGraph." - ] - }, - { - "cell_type": "markdown", - "id": "1abdedbd-d81b-4ee9-b46f-f29439ed1350", - "metadata": {}, - "source": [ - "# Part 1: Planner\n", - "\n", - "\n", - "Largely adapted from [the original source code](https://github.com/SqueezeAILab/LLMCompiler/blob/main/src/llm_compiler/output_parser.py)." - ] - }, - { - "cell_type": "markdown", - "id": "278f76b0-a2e1-42dc-bd3e-6f624984e3dd", - "metadata": {}, - "source": [ - "#### Output Parser\n", - "\n", - "Parses task lists in the following form:\n", - "\n", - "```plaintext\n", - "1. tool_1(\"arg1\", 3.5, ...)\n", - "Thought: I then want to find out Y by using tool_2\n", - "2. tool_2(\"\", ${1})'\n", - "3. join()\"\n", - "```\n", - "\n", - "The \"Thought\" lines are optional. The `${#}` placeholders are variables. These are used to route tool (task) outputs to other tools." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "cc71b6c4-d701-4217-9701-95ac0857e4e1", - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "import re\n", - "from typing import Any, Dict, List, Optional, Sequence, Union\n", - "\n", - "from langchain.agents.agent import AgentOutputParser\n", - "from langchain.schema import OutputParserException\n", - "from langchain_core.tools import BaseTool\n", - "\n", - "THOUGHT_PATTERN = r\"Thought: ([^\\n]*)\"\n", - "# $1 or ${1} -> 1\n", - "ID_PATTERN = r\"\\$\\{?(\\d+)\\}?\"\n", - "END_OF_PLAN = \"\"\n", - "\n", - "\n", - "class ActionParserFSM:\n", - " def __init__(self):\n", - " self.reset()\n", - "\n", - " def reset(self):\n", - " self.state = \"START\"\n", - " self.task_index = \"\"\n", - " self.action = \"\"\n", - " self.comment = \"\"\n", - " self.bracket_count = 0\n", - " self.actions = []\n", - "\n", - " def parse(self, text: str):\n", - " for char in text:\n", - " action = self.process_char(char)\n", - " if action:\n", - " yield action\n", - " action = self.save_action()\n", - " if action:\n", - " yield action\n", - "\n", - " def process_char(self, char: str) -> Optional[dict]:\n", - " action = None\n", - " if self.state == \"START\":\n", - " if char.isdigit():\n", - " self.state = \"NUMBER\"\n", - " self.task_index += char\n", - " elif char == \"\\n\":\n", - " self.reset()\n", - " elif self.state == \"NUMBER\":\n", - " if char == \".\":\n", - " self.state = \"ACTION\"\n", - " elif char.isdigit():\n", - " self.task_index += char\n", - " else:\n", - " self.reset()\n", - " elif self.state == \"ACTION\":\n", - " if char == \"{\":\n", - " self.bracket_count += 1\n", - " elif char == \"}\":\n", - " self.bracket_count -= 1\n", - " if self.bracket_count == 0:\n", - " self.state = \"COMMENT\"\n", - " self.action += char\n", - " elif self.state == \"COMMENT\":\n", - " if char == \"\\n\":\n", - " action = self.save_action()\n", - " self.reset()\n", - " else:\n", - " self.comment += char\n", - " return action\n", - "\n", - " def save_action(self):\n", - " if self.task_index and self.action:\n", - " parsed_action = json.loads(self.action.strip())\n", - " tool_name, args = next(iter(parsed_action.items()))\n", - " return {\n", - " \"task_index\": int(self.task_index),\n", - " \"tool_name\": tool_name,\n", - " \"args\": args,\n", - " }\n", - "\n", - "\n", - "class LLMCompilerPlanParser(AgentOutputParser, extra=\"allow\"):\n", - " \"\"\"Planning output parser.\"\"\"\n", - "\n", - " def __init__(self, tools: Sequence[BaseTool], **kwargs):\n", - " super().__init__(**kwargs)\n", - " self.tools = tools\n", - "\n", - " def parse(self, text: str) -> list[str]:\n", - " parser = ActionParserFSM()\n", - " graph_dict = {}\n", - " for task in parser.parse(text):\n", - " idx = int(task[\"task_index\"])\n", - "\n", - " task = instantiate_task(\n", - " tools=self.tools,\n", - " idx=idx,\n", - " tool_name=task[\"tool_name\"],\n", - " args=task[\"args\"],\n", - " )\n", - "\n", - " graph_dict[idx] = task\n", - " if task[\"tool\"] == \"join\":\n", - " break\n", - "\n", - " return graph_dict\n", - "\n", - "\n", - "### Helper functions\n", - "\n", - "\n", - "def default_dependency_rule(idx, args: str):\n", - " matches = re.findall(ID_PATTERN, args)\n", - " numbers = [int(match) for match in matches]\n", - " return idx in numbers\n", - "\n", - "\n", - "def _get_dependencies_from_graph(\n", - " idx: int, tool_name: str, args: Sequence[Any]\n", - ") -> dict[str, list[str]]:\n", - " \"\"\"Get dependencies from a graph.\"\"\"\n", - " if tool_name == \"join\":\n", - " return list(range(1, idx))\n", - " return [i for i in range(1, idx) if default_dependency_rule(i, str(args))]\n", - "\n", - "\n", - "def instantiate_task(\n", - " tools: Sequence[BaseTool],\n", - " idx: int,\n", - " tool_name: str,\n", - " args: Union[dict, str, bool, None],\n", - ") -> dict:\n", - " dependencies = _get_dependencies_from_graph(idx, tool_name, args)\n", - " if tool_name == \"join\":\n", - " tool = \"join\"\n", - " else:\n", - " try:\n", - " tool = tools[[tool.name for tool in tools].index(tool_name)]\n", - " except ValueError as e:\n", - " raise OutputParserException(f\"Tool {tool_name} not found.\")\n", - " return dict(\n", - " tool=tool,\n", - " args=args,\n", - " dependencies=dependencies,\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "228a2f75-68b1-4dbf-95dc-4bfc738c3b3b", - "metadata": {}, - "source": [ - "#### Planner Code\n", - "\n", - "This takes the input and outputs a plan." - ] - }, - { - "cell_type": "code", - "execution_count": 141, - "id": "58af9746-1011-41e9-a77a-be9cef202aea", - "metadata": {}, - "outputs": [], - "source": [ - "from langchain.chat_models.base import BaseChatModel\n", - "from langchain_core.prompts import ChatPromptTemplate\n", - "from langchain_core.runnables import RunnableBranch\n", - "from langchain_core.tools import BaseTool\n", - "\n", - "END_OF_PLAN = \"\"\n", - "\n", - "JOINER_FINISH = \"Finish\"\n", - "JOINER_REPLAN = \"Replan\"\n", - "\n", - "\n", - "JOIN_DESCRIPTION = (\n", - " \"join():\\n\"\n", - " \" - Collects and combines results from prior actions.\\n\"\n", - " \" - A LLM agent is called upon invoking join to either finalize the user query or wait until the plans are executed.\\n\"\n", - " \" - join should always be the last action in the plan, and will be called in two scenarios:\\n\"\n", - " \" (a) if the answer can be determined by gathering the outputs from tasks to generate the final response.\\n\"\n", - " \" (b) if the answer cannot be determined in the planning phase before you execute the plans. \"\n", - ")\n", - "\n", - "planner_prompt_tmpl_str = (\n", - " \"Given a user query, create a plan to solve it with the utmost parallelizability. \"\n", - " \"Each plan should comprise an action from the following {num_tools} types:\\n\"\n", - " \"{tool_descriptions}\"\n", - " f\"\\n{{num_toolsp1}}. {JOIN_DESCRIPTION}\"\n", - " \"Guidelines:\\n\"\n", - " \" - Each action described above contains input/output types and description.\\n\"\n", - " \" - You must strictly adhere to the input and output types for each action.\\n\"\n", - " \" - The action descriptions contain the guidelines. You MUST strictly follow those guidelines when you use the actions.\\n\"\n", - " \" - Each action in the plan should strictly be one of the above types.\\n\"\n", - " \" - Provide actions ONLY in json form, with the single key being the action name and the value being its arguments. Do not write python code. \\n\"\n", - " \" - Each action line must start with a unique ID, which is strictly increasing.\\n\"\n", - " \" - Inputs for actions can either be constants or outputs from preceding actions. \"\n", - " \"In the latter case, use the format $id to denote the ID of the previous action whose output will be the input.\\n\"\n", - " f\" - Always call join as the last action in the plan. Say '{END_OF_PLAN}' after you call join\\n\"\n", - " \" - Ensure the plan maximizes parallelizability.\\n\"\n", - " \" - Only use the provided action types. If a query cannot be addressed using these, invoke the join action for the next steps.\\n\"\n", - " \" - Never introduce new actions other than the ones provided.\\n\\n\"\n", - " \"{replan}\"\n", - " \"{examples}\"\n", - ")\n", - "\n", - "\n", - "def _generate_planner_prompt(\n", - " tools: Sequence[BaseTool],\n", - " example_prompt=str,\n", - "):\n", - " tool_descriptions = \"\\n\".join(\n", - " f\"{i+1}. {tool.name}: {tool.description}\\n\\tInput schema: {tool.args}\"\n", - " for i, tool in enumerate(tools)\n", - " )\n", - " planner_prompt_template = ChatPromptTemplate.from_messages(\n", - " [(\"system\", planner_prompt_tmpl_str), (\"user\", \"Question: {input}{context}\")]\n", - " ).partial(\n", - " tool_descriptions=tool_descriptions,\n", - " examples=\"Here are some examples:\\n\\n\" + example_prompt\n", - " if example_prompt\n", - " else \"\",\n", - " num_tools=len(tools),\n", - " num_toolsp1=len(tools) + 1,\n", - " )\n", - "\n", - " return planner_prompt_template\n", - "\n", - "\n", - "def create_planner(\n", - " llm: BaseChatModel,\n", - " example_prompt: str,\n", - " tools: Sequence[BaseTool],\n", - " stop: Optional[list[str]] = None,\n", - "):\n", - " og_planner_prompt = _generate_planner_prompt(tools, example_prompt).partial(\n", - " replan=\"\",\n", - " context=\"\",\n", - " )\n", - " replanner_prompt = _generate_planner_prompt(tools, example_prompt).partial(\n", - " replan=' - You are given \"Previous Plan\" which is the plan that the previous agent created along with the execution results '\n", - " \"(given as Observation) of each plan and a general thought (given as Thought) about the executed results.\"\n", - " 'You MUST use these information to create the next plan under \"Current Plan\".\\n'\n", - " ' - When starting the Current Plan, you should start with \"Thought\" that outlines the strategy for the next plan.\\n'\n", - " \" - In the Current Plan, you should NEVER repeat the actions that are already executed in the Previous Plan.\\n\"\n", - " \" - You must continue the task index from the end of the previous one. Do not repeat task indices.\"\n", - " )\n", - " bound_llm = llm.bind(stop=stop)\n", - " return (\n", - " RunnableBranch(\n", - " ((lambda x: x.get(\"context\") is not None), replanner_prompt),\n", - " og_planner_prompt,\n", - " )\n", - " | bound_llm\n", - " | LLMCompilerPlanParser(tools=tools)\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "7feb5c82-b1a9-40ae-863a-0362fe3ce5ea", - "metadata": {}, - "source": [ - "#### Example usage\n", - "\n", - "Here's an example usage of the planner module" - ] - }, - { - "cell_type": "code", - "execution_count": 131, - "id": "3fe074ea-7314-47a7-9a9f-a8e6191ea1f3", - "metadata": {}, - "outputs": [], - "source": [ - "from langchain.tools import tool\n", - "from langchain_openai import ChatOpenAI\n", - "\n", - "\n", - "@tool\n", - "def get_user_id(first_name: str, last_name: Optional[str] = None):\n", - " \"\"\"Query the user IDs of everyone with the provided name.\"\"\"\n", - " return 4\n", - "\n", - "\n", - "@tool\n", - "def get_scores(class_name: str, user_id: int):\n", - " \"\"\"Query the class registry for grades of the provided user ID.\"\"\"\n", - " return \"A+\"\n", - "\n", - "\n", - "examples = (\n", - " \"Question: What's the user ID for Johnny Drop Tables?\\n\"\n", - " '1. {\"get_user_id\": {\"first_name\": \"Johnny\", \"last_name\":\"Drop Tables\"}}\\n'\n", - " f'2. {{\"join\": null}}{END_OF_PLAN}\\n'\n", - " \"###\\n\"\n", - " \"\\n\"\n", - " \"Question: What was Eric Zhang's score in Calc?\\n\"\n", - " '1. {\"get_user_id\": {\"first_name\": \"Eric\", \"last_name\":\"Zhang\"}}\\n'\n", - " '2. {\"get_scores\": {\"class_name\": \"calc\", \"user_id\": \"$1\"}}\\n'\n", - " f'3. {{\"join\": null}}{END_OF_PLAN}\\n'\n", - " \"###\\n\"\n", - " \"\\n\"\n", - ")\n", - "\n", - "planner = create_planner(\n", - " ChatOpenAI(model=\"gpt-3.5-turbo\"),\n", - " example_prompt=examples,\n", - " tools=[get_user_id, get_scores],\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 132, - "id": "55fb0f99-4e59-4e5a-b687-a90bb4d06d39", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{1: {'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: Optional[str] = None) - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", - " 'args': {'first_name': 'Sam', 'last_name': 'Van Damm'},\n", - " 'dependencies': []},\n", - " 2: {'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: Optional[str] = None) - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", - " 'args': {'first_name': 'Will', 'last_name': 'Van Damm'},\n", - " 'dependencies': []},\n", - " 3: {'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", - " 'args': {'class_name': 'Calc BC', 'user_id': '$1'},\n", - " 'dependencies': [1]},\n", - " 4: {'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", - " 'args': {'class_name': 'Calc BC', 'user_id': '$2'},\n", - " 'dependencies': [2]},\n", - " 5: {'tool': 'join', 'args': None, 'dependencies': [1, 2, 3, 4]}}" - ] - }, - "execution_count": 132, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "tasks = planner.invoke(\n", - " {\"input\": \"What are the Calc BC grades for Sam and Will Van Damm?\"}\n", - ")\n", - "tasks" - ] - }, - { - "cell_type": "markdown", - "id": "5d0e795f-61ff-4553-9823-23e7624ca180", - "metadata": {}, - "source": [ - "## 2. Task Fetching Unit\n", - "\n", - "This component schedules the tasks. In the paper, it's kept separate from the \"executor\", but here we create a single DAG to be executed by LangGraph.\n", - "\n", - "Basic idea is that, given a list of dicts of the form:\n", - "\n", - "```typescript\n", - "{\n", - " tool: BaseTool,\n", - " dependencies: number[],\n", - "}\n", - "```\n", - "\n", - "1. Create a topological sort of the tasks\n", - "2. Execute them on the previous step's output, ensuring to perform variable substitution where appropriate" - ] - }, - { - "cell_type": "code", - "execution_count": 184, - "id": "c70a0e28-43db-4ea2-af48-8a2f310dce83", - "metadata": {}, - "outputs": [], - "source": [ - "import functools\n", - "\n", - "from langchain_core.runnables import (\n", - " RunnableLambda,\n", - " RunnableParallel,\n", - " RunnablePassthrough,\n", - ")\n", - "\n", - "\n", - "def _sort_tasks(data):\n", - " if not data:\n", - " return []\n", - " sorted_tasks = []\n", - " # Remove tasks already completed\n", - " min_idx = min([int(k) for k in data])\n", - " data = {\n", - " int(k): {\n", - " **v,\n", - " \"dependencies\": [dep for dep in v[\"dependencies\"] if dep >= min_idx],\n", - " }\n", - " for k, v in data.items()\n", - " }\n", - " while data:\n", - " no_deps = {k: v for k, v in data.items() if not v[\"dependencies\"]}\n", - " if not no_deps:\n", - " raise ValueError(\"We seem to have run into a circular dependency.\")\n", - "\n", - " sorted_tasks.append(no_deps)\n", - " data = {\n", - " k: {\n", - " **v,\n", - " \"dependencies\": [d for d in v[\"dependencies\"] if d not in no_deps],\n", - " }\n", - " for k, v in data.items()\n", - " if k not in no_deps\n", - " }\n", - " return sorted_tasks\n", - "\n", - "\n", - "def _resolve_arg(x: dict, arg: Union[str, Any]):\n", - " if isinstance(arg, str) and arg.startswith(\"$\"):\n", - " try:\n", - " return x[f\"task_{arg[1:]}\"]\n", - " except:\n", - " if arg.endswith(\".output\"):\n", - " return x[f\"task_{arg[1:-7]}\"]\n", - " raise\n", - "\n", - " else:\n", - " return arg\n", - "\n", - "\n", - "def _execute_task(x, task):\n", - " tool_to_use = task[\"tool\"]\n", - " args = task[\"args\"]\n", - " if isinstance(args, str):\n", - " resolved_args = _resolve_arg(x, args)\n", - " elif isinstance(args, dict):\n", - " resolved_args = {key: _resolve_arg(x, val) for key, val in args.items()}\n", - " else:\n", - " # This will likely fail\n", - " resolved_args = args\n", - " try:\n", - " return tool_to_use.invoke(resolved_args)\n", - " except Exception as e:\n", - " return (\n", - " f\"ERROR(Failed to call tool {tool_to_use} with args {tool_to_use}.\"\n", - " + f\" Args resolved to {resolved_args}. Error: {repr(e)})\"\n", - " )\n", - "\n", - "\n", - "def construct_dag(tasks):\n", - " sorted_tasks = _sort_tasks(tasks)\n", - " chain = None\n", - " for idx, task_group in enumerate(sorted_tasks):\n", - " if len(task_group) == 1 and next(iter(task_group.values()))[\"tool\"] == \"join\":\n", - " # TODO: actually join the values\n", - " step = lambda x: {\"join\": x}\n", - " else:\n", - " # Cascade all results forward\n", - " constructor = (\n", - " RunnableParallel if chain is None else RunnablePassthrough.assign\n", - " )\n", - " task_dict = {}\n", - " for idx, task in task_group.items():\n", - " task_dict[f\"task_{idx}\"] = RunnableLambda(\n", - " functools.partial(_execute_task, task=task)\n", - " ).with_config(run_name=f\"task_{idx}\")\n", - "\n", - " step = constructor(**task_dict).with_config(run_name=f\"TaskGroup{idx}\")\n", - " if chain is None:\n", - " chain = step\n", - " else:\n", - " chain |= step\n", - "\n", - " if chain is not None:\n", - " return chain | RunnablePassthrough.assign(tasks=lambda _: tasks)\n", - " return chain" - ] - }, - { - "cell_type": "code", - "execution_count": 185, - "id": "87f11cea-3a8d-479c-8a9e-81223dbbc1f5", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " +------------------------------+ \n", - " | ParallelInput | \n", - " +------------------------------+ \n", - " *** *** \n", - " ** ** \n", - " ** ** \n", - " +-------------+ +-------------+ \n", - " | Lambda(...) | | Lambda(...) | \n", - " +-------------+ +-------------+ \n", - " *** *** \n", - " ** ** \n", - " ** ** \n", - " +-------------------------------+ \n", - " | ParallelOutput | \n", - " +-------------------------------+ \n", - " * \n", - " * \n", - " * \n", - " +------------------------------+ \n", - " | ParallelInput | \n", - " +------------------------------+ \n", - " ***** * ***** \n", - " ***** * ***** \n", - " *** * *** \n", - "+-------------+ +-------------+ +-------------+ \n", - "| Lambda(...) | | Lambda(...) | | Passthrough | \n", - "+-------------+***** +-------------+ *****+-------------+ \n", - " ***** * ***** \n", - " ***** * ***** \n", - " *** * *** \n", - " +-------------------------------+ \n", - " | ParallelOutput | \n", - " +-------------------------------+ \n", - " * \n", - " * \n", - " * \n", - " +-------------------------------+ \n", - " | Lambda(lambda x: {'join': x}) | \n", - " +-------------------------------+ \n", - " * \n", - " * \n", - " * \n", - " +----------------------+ \n", - " | ParallelInput | \n", - " +----------------------+ \n", - " *** *** \n", - " *** *** \n", - " ** ** \n", - " +-------------------------+ +-------------+ \n", - " | Lambda(lambda _: tasks) | | Passthrough | \n", - " +-------------------------+ +-------------+ \n", - " *** *** \n", - " *** *** \n", - " ** ** \n", - " +-----------------------+ \n", - " | ParallelOutput | \n", - " +-----------------------+ \n" - ] - } - ], - "source": [ - "graph = construct_dag(tasks)\n", - "graph.get_graph().print_ascii()" - ] - }, - { - "cell_type": "markdown", - "id": "9efa15ae-817a-48c6-86ed-16bc112fedc5", - "metadata": {}, - "source": [ - "#### Example Plan\n", - "\n", - "We still haven't introduced any cycles in our computation graph, so this is all easily expressed in LCEL." - ] - }, - { - "cell_type": "code", - "execution_count": 135, - "id": "133d124a-0d41-4c6d-a86a-34fa4cb1430f", - "metadata": {}, - "outputs": [], - "source": [ - "chain = planner | construct_dag" - ] - }, - { - "cell_type": "code", - "execution_count": 136, - "id": "55142257-2674-4a47-988e-0d2810917329", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'task_1': 4, 'task_2': 4, 'task_3': 'A+', 'task_4': 'A+'}" - ] - }, - "execution_count": 136, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "example_question = \"Did Aliya get a better score than Roger in Geology?\"\n", - "task_results = chain.invoke({\"input\": example_question})\n", - "task_results[\"join\"]" - ] - }, - { - "cell_type": "markdown", - "id": "563d5311-55f0-4ca1-afbd-01fd970cf3e3", - "metadata": {}, - "source": [ - "## Agent Logic\n", - "\n", - "So now we have the planning and initial execution done. We need a component to process these outputs and either:\n", - "1. Respond with the correct answer.\n", - "2. Loop with a new plan.\n", - "\n", - "The paper calls this the \"joiner\"." - ] - }, - { - "cell_type": "code", - "execution_count": 216, - "id": "2978991e-45a4-44e6-9deb-f941f44fe93a", - "metadata": {}, - "outputs": [], - "source": [ - "from langchain_core.output_parsers import StrOutputParser\n", - "from typing_extensions import TypedDict\n", - "\n", - "\n", - "def format_task(task, idx):\n", - " tool = task[\"tool\"]\n", - " tool_name = tool if isinstance(tool, str) else tool.name # Handle join()\n", - " return f\"{idx}. {{{tool_name}: {task['args']}}}\"\n", - "\n", - "\n", - "def format_tasks(executor_output: dict):\n", - " tasks = executor_output[\"tasks\"]\n", - " prior_observations = executor_output.get(\"observations\")\n", - " formatted_plan = \"\\n\".join(format_task(task, idx) for idx, task in tasks.items())\n", - " observations = \"\\n\".join(f\"{k}: {v}\" for k, v in executor_output[\"join\"].items())\n", - " result = f\"Original Plan:\\n{formatted_plan}\\nExecuted plan results:\\n{observations}\"\n", - " if prior_observations:\n", - " result += f\"\\nPrevious Results:\\n{prior_observations}\"\n", - " return result\n", - "\n", - "\n", - "def _parse_joiner_output(raw_answer: str) -> str:\n", - " thought, answer, is_replan = \"\", \"\", False # default values\n", - " raw_answers = raw_answer.split(\"\\n\")\n", - " for ans in raw_answers:\n", - " if ans.startswith(\"Action:\"):\n", - " answer = ans[ans.find(\"(\") + 1 : ans.find(\")\")]\n", - " is_replan = JOINER_REPLAN in ans\n", - " elif ans.startswith(\"Thought:\"):\n", - " thought = ans.split(\"Thought:\")[1].strip()\n", - " if is_replan:\n", - " return {\"thought\": thought, \"context\": answer}\n", - " else:\n", - " return {\"thought\": thought, \"answer\": answer}" - ] - }, - { - "cell_type": "code", - "execution_count": 217, - "id": "b2feea5a-e0e4-4cff-8cb5-fdbfec95ba57", - "metadata": {}, - "outputs": [], - "source": [ - "def create_joiner(prompt, llm):\n", - " return (\n", - " (\n", - " lambda x: {\n", - " **x[\"plan\"],\n", - " \"input\": x[\"input\"],\n", - " \"context\": x.get(\"context\"),\n", - " \"observations\": x.get(\"observations\"),\n", - " }\n", - " )\n", - " | RunnablePassthrough.assign(scratchpad=format_tasks)\n", - " | ChatPromptTemplate.from_messages([(\"system\", prompt), (\"user\", \"{input}\")])\n", - " | llm\n", - " | StrOutputParser()\n", - " | _parse_joiner_output\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": 218, - "id": "942dab42-ad42-4ba2-90d5-49edbe4fae68", - "metadata": {}, - "outputs": [], - "source": [ - "system_prompt = (\n", - " \"Solve a question answering task. Here are some guidelines:\\n\"\n", - " \" - In the Assistant Scratchpad, you will be given results of a plan you have executed to answer the user's question.\\n\"\n", - " \" - Thought needs to reason about the question based on the Observations in 1-2 sentences.\\n\"\n", - " \" - Ignore irrelevant action results.\\n\"\n", - " \" - If the required information is present, give a concise but complete and helpful answer to the user's question.\\n\"\n", - " \" - If you are unable to give a satisfactory finishing answer, replan to get the required information.\"\n", - " \" Respond in the following format:\\n\\n\"\n", - " \"Thought: \\n\"\n", - " \"Action: \\n\"\n", - " \"Available actions:\\n\"\n", - " f\" (1) {JOINER_FINISH}(the final answer to return to the user): returns the answer and finishes the task.\\n\"\n", - " f\" (2) {JOINER_REPLAN}(the reasoning and other information that will help you plan again. Can be a line of any length): instructs why we must replan\\n\\n\"\n", - " \" Examples:\\n\"\n", - " \"Question: How many users are currently using the new product?\\n\"\n", - " \"...task returns the number 32,000\\n\"\n", - " \"Thought: I find no issue with the original plan, and the results satisfy everything in the user question.\\n\"\n", - " f\"Action: {JOINER_FINISH}(32,000 users currently use the new product)\\n###\\n\"\n", - " \"Question: How much cooler is it in NY than SF?\\n\"\n", - " \"...task results show SF is 57 degrees fahrenheit today, and they show in NY it has a high of 32 degrees fahrenheit \\n\"\n", - " \"Thought: I can answer by synthesizing the results.\\n\"\n", - " f\"Action: {JOINER_FINISH}(NY is 25 degrees cooler than SF today, as it has a high of 32 degrees Fahrenheit today, whereas in SF, it is 57 degrees Fahrenheit.)\\n###\\n\"\n", - " \"Question: Are the gophers beating the rabbits??\\n\"\n", - " \"...task returns the a score of 7 for rabbits but no other value...\\n\"\n", - " \"Thought: I need the gophers' score to make a final decision.\\n\"\n", - " f\"Action: {JOINER_REPLAN}(The rabbits have a score of 7, but I need the gophers' score.)\"\n", - " \"\\nAssistant Scratchpad:\\n{scratchpad}\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 219, - "id": "a07d0804-2ce5-4462-98eb-4473f36ef704", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{'thought': 'The plan has been executed successfully and returned the scores for both Aliya and Roger in Geology. We can compare their scores to determine if Aliya got a better score than Roger.',\n", - " 'answer': 'Aliya got an A+ in Geology, while Roger also got an A+. Therefore, they both got the same score in Geology.'}" - ] - }, - "execution_count": 219, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "joiner = create_joiner(system_prompt, ChatOpenAI(model=\"gpt-3.5-turbo\"))\n", - "joiner.invoke({\"plan\": task_results, \"input\": example_question})" - ] - }, - { - "cell_type": "markdown", - "id": "b099e5ee-2c23-47d9-9387-0f64e02627d3", - "metadata": {}, - "source": [ - "### Construct Agent\n", - "\n", - "Now we have all the required pieces! Let's construct our agent with its tools." - ] - }, - { - "cell_type": "code", - "execution_count": 220, - "id": "565b08d2-d19c-4125-97f7-996fc01bc631", - "metadata": {}, - "outputs": [], - "source": [ - "import getpass\n", - "import os\n", - "\n", - "os.environ[\"TAVILY_API_KEY\"] = (\n", - " os.environ.get(\"TAVILY_API_KEY\")\n", - " if \"TAVILY_API_KEY\" in os.environ\n", - " else getpass.getpass(\"Tavily API Key:\")\n", - ")\n", - "# Then fetch a credentials.json file\n", - "# https://developers.google.com/gmail/api/quickstart/python#authorize_credentials_for_a_desktop_application" - ] - }, - { - "cell_type": "code", - "execution_count": 221, - "id": "bd84ca4b-eacb-471d-9975-5449047b5bed", - "metadata": {}, - "outputs": [], - "source": [ - "from operator import add, mul, sub, truediv\n", - "from typing import Literal\n", - "\n", - "from langchain_community.agent_toolkits import GmailToolkit\n", - "from langchain_community.tools.tavily_search import TavilySearchResults\n", - "from langchain_core.tools import tool\n", - "\n", - "\n", - "@tool\n", - "def calculate(\n", - " arg1: float,\n", - " arg2: float,\n", - " op: Union[Literal[\"+\"], Literal[\"-\"], Literal[\"*\"], Literal[\"/\"]],\n", - "):\n", - " \"\"\"Calculate a mathematical operation on two arguments.\"\"\"\n", - " resolved_op = {\"+\": add, \"-\": sub, \"*\": mul, \"/\": truediv}\n", - " return resolved_op[op](arg1, arg2)\n", - "\n", - "\n", - "tools = [TavilySearchResults(max_results=1), calculate]" - ] - }, - { - "cell_type": "code", - "execution_count": 224, - "id": "768b5f11-e3d2-47be-8143-a7dcd8765243", - "metadata": {}, - "outputs": [], - "source": [ - "from langgraph.graph import END, StateGraph\n", - "\n", - "\n", - "class GraphState(TypedDict):\n", - " input: str\n", - " plan: Dict\n", - " agent_output: Dict\n", - " observations: Dict\n", - " num_iterations: int\n", - " context: str\n", - " stop_reason: str\n", - "\n", - "\n", - "MAX_ITERATIONS = 5\n", - "workflow = StateGraph(GraphState)\n", - "\n", - "# 1. Define vertices\n", - "\n", - "planner = create_planner(\n", - " llm=ChatOpenAI(model=\"gpt-4-1106-preview\"),\n", - " # Add more examples to improve reliability\n", - " example_prompt=(\n", - " \"Question: What's the capital of Myanmar?\\n\"\n", - " '1. {\"tavily_search_results_json\": {\"query\": \"Capital of Myanmar\"}}\\n'\n", - " f'2. {{\"join\": null}}{END_OF_PLAN}\\n'\n", - " \"###\\n\"\n", - " \"\\n\"\n", - " ),\n", - " tools=tools,\n", - ")\n", - "\n", - "plan_and_execute = planner | construct_dag\n", - "joiner = create_joiner(system_prompt, ChatOpenAI(model=\"gpt-4-1106-preview\"))\n", - "\n", - "\n", - "def _reformat_task(idx, task: Union[BaseTool, str]):\n", - " tool = task[\"tool\"]\n", - " tool_name = tool if isinstance(tool, str) else tool.name\n", - " called = {tool_name: task[\"args\"]}\n", - " return f\"{idx}. {json.dumps(called)}\"\n", - "\n", - "\n", - "def provide_context(state):\n", - " # Insert a context string for the re-planner.\n", - " # This could alternatively call an LLM to provide additional logic\n", - " context = state[\"agent_output\"][\"context\"]\n", - " num_iterations = int(state.get(\"num_iterations\") or 1) + 1\n", - " previous_plan = \"\\n\".join(\n", - " [\n", - " _reformat_task(idx, task)\n", - " for idx, task in sorted(state[\"plan\"][\"tasks\"].items())\n", - " ]\n", - " )\n", - " context_str = (\n", - " f\"\\n\\nPrevious Plan:\\n{previous_plan}\\n\"\n", - " f\"{context}\\nYou have made {num_iterations}/{MAX_ITERATIONS} attempts thus far.\"\n", - " )\n", - " observations = state[\"observations\"] or {}\n", - " for task, observation in state[\"plan\"][\"join\"].items():\n", - " observations[task] = observation\n", - " return {\n", - " \"context\": context_str,\n", - " \"num_iterations\": num_iterations,\n", - " \"observations\": observations,\n", - " }\n", - "\n", - "\n", - "def add_stop_reason(state):\n", - " num_iterations = int(state.get(\"num_iterations\") or 0)\n", - " if num_iterations >= MAX_ITERATIONS:\n", - " return {\"stop_reason\": \"end_max_iter\"}\n", - " if state[\"agent_output\"].get(\"answer\"):\n", - " return {\"stop_reason\": \"answer\"}\n", - " return {\"stop_reason\": None}\n", - "\n", - "\n", - "# Assign each node to a state variable to update\n", - "workflow.add_node(\"plan_and_execute\", RunnablePassthrough.assign(plan=plan_and_execute))\n", - "workflow.add_node(\"join\", RunnablePassthrough.assign(agent_output=joiner))\n", - "workflow.add_node(\"provide_context\", provide_context)\n", - "workflow.add_node(\"provide_stop_reason\", add_stop_reason)\n", - "\n", - "\n", - "## Define edges\n", - "\n", - "workflow.add_edge(\"plan_and_execute\", \"join\")\n", - "workflow.add_edge(\"provide_context\", \"plan_and_execute\")\n", - "workflow.add_edge(\"join\", \"provide_stop_reason\")\n", - "\n", - "### This condition determines looping logic\n", - "\n", - "\n", - "def should_continue(state):\n", - " if state[\"stop_reason\"] is None:\n", - " return \"continue\"\n", - " return \"end\"\n", - "\n", - "\n", - "workflow.add_conditional_edges(\n", - " start_key=\"provide_stop_reason\",\n", - " # Next, we pass in the function that will determine which node is called next.\n", - " condition=should_continue,\n", - " conditional_edge_mapping={\n", - " # If it generates context, we must replan\n", - " \"continue\": \"provide_context\",\n", - " # Otherwise we finish.\n", - " \"end\": END,\n", - " },\n", - ")\n", - "workflow.set_entry_point(\"plan_and_execute\")\n", - "chain = workflow.compile()" - ] - }, - { - "cell_type": "markdown", - "id": "9f8c9849-8531-463d-a0ef-dcc3d9888b2d", - "metadata": {}, - "source": [ - "## Simple question\n", - "\n", - "Let's ask a simple question of the agent." - ] - }, - { - "cell_type": "code", - "execution_count": 225, - "id": "5bc4584a-e31c-4065-805e-76a6db30676a", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "The GDP of New York in 2022 was about 1.56 trillion U.S. dollars.\n" - ] - } - ], - "source": [ - "result = chain.invoke({\"input\": \"What's the GDP of New York?\"})\n", - "print(result[\"agent_output\"][\"answer\"])" - ] - }, - { - "cell_type": "markdown", - "id": "33c65ef5-b4b2-4ab2-8c78-a551da7819b9", - "metadata": {}, - "source": [ - "## Multi-hop question" - ] - }, - { - "cell_type": "code", - "execution_count": 227, - "id": "0b3a0916-d8ca-4092-b91c-d9e2b05259d8", - "metadata": {}, - "outputs": [ - { - "ename": "JSONDecodeError", - "evalue": "Expecting value: line 1 column 1 (char 0)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mJSONDecodeError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[227], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mchain\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[43m \u001b[49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minput\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mHow much larger is the GDP of the UK than that of New York?\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:492\u001b[0m, in \u001b[0;36mPregel.invoke\u001b[0;34m(self, input, config, output_keys, input_keys, **kwargs)\u001b[0m\n\u001b[1;32m 482\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m 483\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 484\u001b[0m \u001b[38;5;28minput\u001b[39m: Union[\u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], Any],\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 489\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 490\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[\u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], Any]:\n\u001b[1;32m 491\u001b[0m latest: Union[\u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], Any] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 492\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstream\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 493\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 494\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 495\u001b[0m \u001b[43m \u001b[49m\u001b[43moutput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_keys\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43moutput_keys\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 496\u001b[0m \u001b[43m \u001b[49m\u001b[43minput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 497\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 498\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 499\u001b[0m \u001b[43m \u001b[49m\u001b[43mlatest\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\n\u001b[1;32m 500\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m latest\n", - "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:528\u001b[0m, in \u001b[0;36mPregel.transform\u001b[0;34m(self, input, config, output_keys, input_keys, **kwargs)\u001b[0m\n\u001b[1;32m 519\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtransform\u001b[39m(\n\u001b[1;32m 520\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 521\u001b[0m \u001b[38;5;28minput\u001b[39m: Iterator[Union[\u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], Any]],\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 526\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 527\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Iterator[Union[\u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], Any]]:\n\u001b[0;32m--> 528\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_transform_stream_with_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 529\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 530\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_transform\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 531\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 532\u001b[0m \u001b[43m \u001b[49m\u001b[43moutput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 533\u001b[0m \u001b[43m \u001b[49m\u001b[43minput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 534\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 535\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 536\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01myield\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1226\u001b[0m, in \u001b[0;36mRunnable._transform_stream_with_config\u001b[0;34m(self, input, transformer, config, run_type, **kwargs)\u001b[0m\n\u001b[1;32m 1224\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1225\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m-> 1226\u001b[0m chunk: Output \u001b[38;5;241m=\u001b[39m context\u001b[38;5;241m.\u001b[39mrun(\u001b[38;5;28mnext\u001b[39m, iterator) \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m 1227\u001b[0m \u001b[38;5;28;01myield\u001b[39;00m chunk\n\u001b[1;32m 1228\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m final_output_supported:\n", - "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:313\u001b[0m, in \u001b[0;36mPregel._transform\u001b[0;34m(self, input, run_manager, config, input_keys, output_keys)\u001b[0m\n\u001b[1;32m 303\u001b[0m done, inflight \u001b[38;5;241m=\u001b[39m concurrent\u001b[38;5;241m.\u001b[39mfutures\u001b[38;5;241m.\u001b[39mwait(\n\u001b[1;32m 304\u001b[0m [\n\u001b[1;32m 305\u001b[0m executor\u001b[38;5;241m.\u001b[39msubmit(proc\u001b[38;5;241m.\u001b[39minvoke, \u001b[38;5;28minput\u001b[39m, config)\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 309\u001b[0m timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstep_timeout,\n\u001b[1;32m 310\u001b[0m )\n\u001b[1;32m 312\u001b[0m \u001b[38;5;66;03m# interrupt on failure or timeout\u001b[39;00m\n\u001b[0;32m--> 313\u001b[0m \u001b[43m_interrupt_or_proceed\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdone\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minflight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstep\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 315\u001b[0m \u001b[38;5;66;03m# apply writes to channels\u001b[39;00m\n\u001b[1;32m 316\u001b[0m _apply_writes(checkpoint, channels, pending_writes, config, step \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m)\n", - "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:611\u001b[0m, in \u001b[0;36m_interrupt_or_proceed\u001b[0;34m(done, inflight, step)\u001b[0m\n\u001b[1;32m 609\u001b[0m inflight\u001b[38;5;241m.\u001b[39mpop()\u001b[38;5;241m.\u001b[39mcancel()\n\u001b[1;32m 610\u001b[0m \u001b[38;5;66;03m# raise the exception\u001b[39;00m\n\u001b[0;32m--> 611\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc\n\u001b[1;32m 612\u001b[0m \u001b[38;5;66;03m# TODO this is where retry of an entire step would happen\u001b[39;00m\n\u001b[1;32m 614\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inflight:\n\u001b[1;32m 615\u001b[0m \u001b[38;5;66;03m# if we got here means we timed out\u001b[39;00m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/thread.py:58\u001b[0m, in \u001b[0;36m_WorkItem.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m 57\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 58\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 59\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m 60\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39mset_exception(exc)\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:3596\u001b[0m, in \u001b[0;36mRunnableBindingBase.invoke\u001b[0;34m(self, input, config, **kwargs)\u001b[0m\n\u001b[1;32m 3590\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m 3591\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 3592\u001b[0m \u001b[38;5;28minput\u001b[39m: Input,\n\u001b[1;32m 3593\u001b[0m config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m 3594\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Optional[Any],\n\u001b[1;32m 3595\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Output:\n\u001b[0;32m-> 3596\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbound\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 3597\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3598\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_merge_configs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3599\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3600\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1774\u001b[0m, in \u001b[0;36mRunnableSequence.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m 1772\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1773\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msteps):\n\u001b[0;32m-> 1774\u001b[0m \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1775\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1776\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;66;43;03m# mark each step as a child run\u001b[39;49;00m\n\u001b[1;32m 1777\u001b[0m \u001b[43m \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1778\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseq:step:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mi\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1779\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1780\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1781\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m 1782\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/passthrough.py:415\u001b[0m, in \u001b[0;36mRunnableAssign.invoke\u001b[0;34m(self, input, config, **kwargs)\u001b[0m\n\u001b[1;32m 409\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m 410\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 411\u001b[0m \u001b[38;5;28minput\u001b[39m: Dict[\u001b[38;5;28mstr\u001b[39m, Any],\n\u001b[1;32m 412\u001b[0m config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m 413\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 414\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Dict[\u001b[38;5;28mstr\u001b[39m, Any]:\n\u001b[0;32m--> 415\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_with_config\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_invoke\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:975\u001b[0m, in \u001b[0;36mRunnable._call_with_config\u001b[0;34m(self, func, input, config, run_type, **kwargs)\u001b[0m\n\u001b[1;32m 971\u001b[0m context \u001b[38;5;241m=\u001b[39m copy_context()\n\u001b[1;32m 972\u001b[0m context\u001b[38;5;241m.\u001b[39mrun(var_child_runnable_config\u001b[38;5;241m.\u001b[39mset, child_config)\n\u001b[1;32m 973\u001b[0m output \u001b[38;5;241m=\u001b[39m cast(\n\u001b[1;32m 974\u001b[0m Output,\n\u001b[0;32m--> 975\u001b[0m \u001b[43mcontext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 976\u001b[0m \u001b[43m \u001b[49m\u001b[43mcall_func_with_variable_args\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 977\u001b[0m \u001b[43m \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m 978\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m 979\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 980\u001b[0m \u001b[43m \u001b[49m\u001b[43mrun_manager\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 981\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 982\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 983\u001b[0m )\n\u001b[1;32m 984\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m 985\u001b[0m run_manager\u001b[38;5;241m.\u001b[39mon_chain_error(e)\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/config.py:323\u001b[0m, in \u001b[0;36mcall_func_with_variable_args\u001b[0;34m(func, input, config, run_manager, **kwargs)\u001b[0m\n\u001b[1;32m 321\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m run_manager \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m accepts_run_manager(func):\n\u001b[1;32m 322\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_manager\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m run_manager\n\u001b[0;32m--> 323\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/passthrough.py:402\u001b[0m, in \u001b[0;36mRunnableAssign._invoke\u001b[0;34m(self, input, run_manager, config, **kwargs)\u001b[0m\n\u001b[1;32m 389\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_invoke\u001b[39m(\n\u001b[1;32m 390\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 391\u001b[0m \u001b[38;5;28minput\u001b[39m: Dict[\u001b[38;5;28mstr\u001b[39m, Any],\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 394\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m 395\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Dict[\u001b[38;5;28mstr\u001b[39m, Any]:\n\u001b[1;32m 396\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\n\u001b[1;32m 397\u001b[0m \u001b[38;5;28minput\u001b[39m, \u001b[38;5;28mdict\u001b[39m\n\u001b[1;32m 398\u001b[0m ), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe input to RunnablePassthrough.assign() must be a dict.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 400\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m {\n\u001b[1;32m 401\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;28minput\u001b[39m,\n\u001b[0;32m--> 402\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmapper\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 403\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 404\u001b[0m \u001b[43m \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 405\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 406\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 407\u001b[0m }\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:2339\u001b[0m, in \u001b[0;36mRunnableParallel.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m 2326\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m get_executor_for_config(config) \u001b[38;5;28;01mas\u001b[39;00m executor:\n\u001b[1;32m 2327\u001b[0m futures \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 2328\u001b[0m executor\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m 2329\u001b[0m step\u001b[38;5;241m.\u001b[39minvoke,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 2337\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, step \u001b[38;5;129;01min\u001b[39;00m steps\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 2338\u001b[0m ]\n\u001b[0;32m-> 2339\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43m{\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mzip\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43msteps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\u001b[43m)\u001b[49m\u001b[43m}\u001b[49m\n\u001b[1;32m 2340\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m 2341\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:2339\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 2326\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m get_executor_for_config(config) \u001b[38;5;28;01mas\u001b[39;00m executor:\n\u001b[1;32m 2327\u001b[0m futures \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 2328\u001b[0m executor\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m 2329\u001b[0m step\u001b[38;5;241m.\u001b[39minvoke,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 2337\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, step \u001b[38;5;129;01min\u001b[39;00m steps\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 2338\u001b[0m ]\n\u001b[0;32m-> 2339\u001b[0m output \u001b[38;5;241m=\u001b[39m {key: \u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m key, future \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(steps, futures)}\n\u001b[1;32m 2340\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m 2341\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/_base.py:456\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m 454\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m 455\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 456\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 457\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 458\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m()\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m 400\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m 402\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m 403\u001b[0m \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m 404\u001b[0m \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/thread.py:58\u001b[0m, in \u001b[0;36m_WorkItem.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m 57\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 58\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 59\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m 60\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39mset_exception(exc)\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1774\u001b[0m, in \u001b[0;36mRunnableSequence.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m 1772\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1773\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msteps):\n\u001b[0;32m-> 1774\u001b[0m \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1775\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1776\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;66;43;03m# mark each step as a child run\u001b[39;49;00m\n\u001b[1;32m 1777\u001b[0m \u001b[43m \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1778\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseq:step:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mi\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1779\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1780\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1781\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m 1782\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:167\u001b[0m, in \u001b[0;36mBaseOutputParser.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m 164\u001b[0m \u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Union[\u001b[38;5;28mstr\u001b[39m, BaseMessage], config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 165\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m 166\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28minput\u001b[39m, BaseMessage):\n\u001b[0;32m--> 167\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_with_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 168\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43minner_input\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse_result\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 169\u001b[0m \u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43mChatGeneration\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmessage\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minner_input\u001b[49m\u001b[43m)\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 171\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 172\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 173\u001b[0m \u001b[43m \u001b[49m\u001b[43mrun_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mparser\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 174\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_with_config(\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mlambda\u001b[39;00m inner_input: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparse_result([Generation(text\u001b[38;5;241m=\u001b[39minner_input)]),\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28minput\u001b[39m,\n\u001b[1;32m 179\u001b[0m config,\n\u001b[1;32m 180\u001b[0m run_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mparser\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 181\u001b[0m )\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:975\u001b[0m, in \u001b[0;36mRunnable._call_with_config\u001b[0;34m(self, func, input, config, run_type, **kwargs)\u001b[0m\n\u001b[1;32m 971\u001b[0m context \u001b[38;5;241m=\u001b[39m copy_context()\n\u001b[1;32m 972\u001b[0m context\u001b[38;5;241m.\u001b[39mrun(var_child_runnable_config\u001b[38;5;241m.\u001b[39mset, child_config)\n\u001b[1;32m 973\u001b[0m output \u001b[38;5;241m=\u001b[39m cast(\n\u001b[1;32m 974\u001b[0m Output,\n\u001b[0;32m--> 975\u001b[0m \u001b[43mcontext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 976\u001b[0m \u001b[43m \u001b[49m\u001b[43mcall_func_with_variable_args\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 977\u001b[0m \u001b[43m \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m 978\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m 979\u001b[0m \u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 980\u001b[0m \u001b[43m \u001b[49m\u001b[43mrun_manager\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 981\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 982\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 983\u001b[0m )\n\u001b[1;32m 984\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m 985\u001b[0m run_manager\u001b[38;5;241m.\u001b[39mon_chain_error(e)\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/config.py:323\u001b[0m, in \u001b[0;36mcall_func_with_variable_args\u001b[0;34m(func, input, config, run_manager, **kwargs)\u001b[0m\n\u001b[1;32m 321\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m run_manager \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m accepts_run_manager(func):\n\u001b[1;32m 322\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_manager\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m run_manager\n\u001b[0;32m--> 323\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:168\u001b[0m, in \u001b[0;36mBaseOutputParser.invoke..\u001b[0;34m(inner_input)\u001b[0m\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m 164\u001b[0m \u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Union[\u001b[38;5;28mstr\u001b[39m, BaseMessage], config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 165\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m 166\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28minput\u001b[39m, BaseMessage):\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_with_config(\n\u001b[0;32m--> 168\u001b[0m \u001b[38;5;28;01mlambda\u001b[39;00m inner_input: \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse_result\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 169\u001b[0m \u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43mChatGeneration\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmessage\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minner_input\u001b[49m\u001b[43m)\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 171\u001b[0m \u001b[38;5;28minput\u001b[39m,\n\u001b[1;32m 172\u001b[0m config,\n\u001b[1;32m 173\u001b[0m run_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mparser\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 174\u001b[0m )\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_with_config(\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mlambda\u001b[39;00m inner_input: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparse_result([Generation(text\u001b[38;5;241m=\u001b[39minner_input)]),\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28minput\u001b[39m,\n\u001b[1;32m 179\u001b[0m config,\n\u001b[1;32m 180\u001b[0m run_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mparser\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 181\u001b[0m )\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:219\u001b[0m, in \u001b[0;36mBaseOutputParser.parse_result\u001b[0;34m(self, result, partial)\u001b[0m\n\u001b[1;32m 206\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mparse_result\u001b[39m(\u001b[38;5;28mself\u001b[39m, result: List[Generation], \u001b[38;5;241m*\u001b[39m, partial: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m 207\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Parse a list of candidate model Generations into a specific format.\u001b[39;00m\n\u001b[1;32m 208\u001b[0m \n\u001b[1;32m 209\u001b[0m \u001b[38;5;124;03m The return value is parsed from only the first Generation in the result, which\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 217\u001b[0m \u001b[38;5;124;03m Structured output.\u001b[39;00m\n\u001b[1;32m 218\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 219\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtext\u001b[49m\u001b[43m)\u001b[49m\n", - "Cell \u001b[0;32mIn[1], line 88\u001b[0m, in \u001b[0;36mLLMCompilerPlanParser.parse\u001b[0;34m(self, text)\u001b[0m\n\u001b[1;32m 86\u001b[0m parser \u001b[38;5;241m=\u001b[39m ActionParserFSM()\n\u001b[1;32m 87\u001b[0m graph_dict \u001b[38;5;241m=\u001b[39m {}\n\u001b[0;32m---> 88\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtask\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mparser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparse\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtext\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 89\u001b[0m \u001b[43m \u001b[49m\u001b[43midx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtask\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtask_index\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[43m \u001b[49m\u001b[43mtask\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43minstantiate_task\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 92\u001b[0m \u001b[43m \u001b[49m\u001b[43mtools\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtools\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 93\u001b[0m \u001b[43m \u001b[49m\u001b[43midx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43midx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 94\u001b[0m \u001b[43m \u001b[49m\u001b[43mtool_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtask\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtool_name\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 95\u001b[0m \u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtask\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43margs\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 96\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "Cell \u001b[0;32mIn[1], line 29\u001b[0m, in \u001b[0;36mActionParserFSM.parse\u001b[0;34m(self, text)\u001b[0m\n\u001b[1;32m 27\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mparse\u001b[39m(\u001b[38;5;28mself\u001b[39m, text: \u001b[38;5;28mstr\u001b[39m):\n\u001b[1;32m 28\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m char \u001b[38;5;129;01min\u001b[39;00m text:\n\u001b[0;32m---> 29\u001b[0m action \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_char\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchar\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 30\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m action:\n\u001b[1;32m 31\u001b[0m \u001b[38;5;28;01myield\u001b[39;00m action\n", - "Cell \u001b[0;32mIn[1], line 61\u001b[0m, in \u001b[0;36mActionParserFSM.process_char\u001b[0;34m(self, char)\u001b[0m\n\u001b[1;32m 59\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCOMMENT\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 60\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m char \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 61\u001b[0m action \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave_action\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 62\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset()\n\u001b[1;32m 63\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "Cell \u001b[0;32mIn[1], line 69\u001b[0m, in \u001b[0;36mActionParserFSM.save_action\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 67\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msave_action\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m 68\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask_index \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction:\n\u001b[0;32m---> 69\u001b[0m parsed_action \u001b[38;5;241m=\u001b[39m \u001b[43mjson\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstrip\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 70\u001b[0m tool_name, args \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(parsed_action\u001b[38;5;241m.\u001b[39mitems()))\n\u001b[1;32m 71\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m {\n\u001b[1;32m 72\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtask_index\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mint\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask_index),\n\u001b[1;32m 73\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtool_name\u001b[39m\u001b[38;5;124m\"\u001b[39m: tool_name,\n\u001b[1;32m 74\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124margs\u001b[39m\u001b[38;5;124m\"\u001b[39m: args,\n\u001b[1;32m 75\u001b[0m }\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/json/__init__.py:346\u001b[0m, in \u001b[0;36mloads\u001b[0;34m(s, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)\u001b[0m\n\u001b[1;32m 341\u001b[0m s \u001b[38;5;241m=\u001b[39m s\u001b[38;5;241m.\u001b[39mdecode(detect_encoding(s), \u001b[38;5;124m'\u001b[39m\u001b[38;5;124msurrogatepass\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 343\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\u001b[38;5;28mcls\u001b[39m \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m object_hook \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m\n\u001b[1;32m 344\u001b[0m parse_int \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m parse_float \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m\n\u001b[1;32m 345\u001b[0m parse_constant \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m object_pairs_hook \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m kw):\n\u001b[0;32m--> 346\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_default_decoder\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdecode\u001b[49m\u001b[43m(\u001b[49m\u001b[43ms\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 347\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mcls\u001b[39m \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 348\u001b[0m \u001b[38;5;28mcls\u001b[39m \u001b[38;5;241m=\u001b[39m JSONDecoder\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/json/decoder.py:337\u001b[0m, in \u001b[0;36mJSONDecoder.decode\u001b[0;34m(self, s, _w)\u001b[0m\n\u001b[1;32m 332\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecode\u001b[39m(\u001b[38;5;28mself\u001b[39m, s, _w\u001b[38;5;241m=\u001b[39mWHITESPACE\u001b[38;5;241m.\u001b[39mmatch):\n\u001b[1;32m 333\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return the Python representation of ``s`` (a ``str`` instance\u001b[39;00m\n\u001b[1;32m 334\u001b[0m \u001b[38;5;124;03m containing a JSON document).\u001b[39;00m\n\u001b[1;32m 335\u001b[0m \n\u001b[1;32m 336\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 337\u001b[0m obj, end \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw_decode\u001b[49m\u001b[43m(\u001b[49m\u001b[43ms\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43midx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_w\u001b[49m\u001b[43m(\u001b[49m\u001b[43ms\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mend\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 338\u001b[0m end \u001b[38;5;241m=\u001b[39m _w(s, end)\u001b[38;5;241m.\u001b[39mend()\n\u001b[1;32m 339\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m end \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mlen\u001b[39m(s):\n", - "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/json/decoder.py:355\u001b[0m, in \u001b[0;36mJSONDecoder.raw_decode\u001b[0;34m(self, s, idx)\u001b[0m\n\u001b[1;32m 353\u001b[0m obj, end \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscan_once(s, idx)\n\u001b[1;32m 354\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 355\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m JSONDecodeError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExpecting value\u001b[39m\u001b[38;5;124m\"\u001b[39m, s, err\u001b[38;5;241m.\u001b[39mvalue) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 356\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj, end\n", - "\u001b[0;31mJSONDecodeError\u001b[0m: Expecting value: line 1 column 1 (char 0)" - ] - } - ], - "source": [ - "result = chain.invoke(\n", - " {\"input\": \"How much larger is the GDP of the UK than that of New York?\"}\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a157cd88-2e54-4525-8126-6d22affa31d7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.2" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/examples/advanced_agents/llm_compiler/LLMCompiler.ipynb b/examples/advanced_agents/llm_compiler/LLMCompiler.ipynb new file mode 100644 index 000000000..20e0097fd --- /dev/null +++ b/examples/advanced_agents/llm_compiler/LLMCompiler.ipynb @@ -0,0 +1,1036 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0c8b472b-f3fb-46c2-841f-930a4692697b", + "metadata": {}, + "source": [ + "# Implementing LLMCompiler using LangGraph\n", + "By Kim, et. al [🔗](https://arxiv.org/abs/2312.04511)\n", + "\n", + "LLMCompiler is an agent architecture intented on speeding up the latency of agentic tasks via fast, parallel tool execution. It has 3 main components:\n", + "\n", + "1. Planner: generate a DAG of tasks.\n", + "2. Task Fetching Unit: schedules and executes the tasks\n", + "3. Joiner: Responds to the user or triggers a second plan\n", + "\n", + "\n", + "This notebook walks through each component and shows how to wire them together using LangGraph." + ] + }, + { + "cell_type": "markdown", + "id": "1abdedbd-d81b-4ee9-b46f-f29439ed1350", + "metadata": {}, + "source": [ + "# Part 1: Planner\n", + "\n", + "\n", + "Largely adapted from [the original source code](https://github.com/SqueezeAILab/LLMCompiler/blob/main/src/llm_compiler/output_parser.py), the planner accepts the input question and generates a task list to execute.\n", + "\n", + "If it is provided with a previous plan, it is instructed to re-plan, which is useful if, upon completion of the first batch of tasks, the agent must take more actions.\n", + "\n", + "The code below composes constructs the prompt template for the planner and composes it with LLM and output parser, defined in [output_parser.py](./output_parser.py). The output parser processes a task list in the following form:\n", + "\n", + "```plaintext\n", + "1. tool_1(arg1=\"arg1\", arg2=3.5, ...)\n", + "Thought: I then want to find out Y by using tool_2\n", + "2. tool_2(arg1=\"\", arg2=\"${1}\")'\n", + "3. join()\"\n", + "```\n", + "\n", + "The \"Thought\" lines are optional. The `${#}` placeholders are variables. These are used to route tool (task) outputs to other tools." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "15dd9639-691f-4906-9012-83fd6e9ac126", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Optional, Sequence\n", + "\n", + "from langchain.chat_models.base import BaseChatModel\n", + "from langchain_core.prompts import ChatPromptTemplate\n", + "from langchain_core.runnables import RunnableBranch\n", + "from langchain_core.tools import BaseTool\n", + "from llm_compiler.output_parser import LLMCompilerPlanParser\n", + "\n", + "END_OF_PLAN = \"\"\n", + "\n", + "\n", + "# The required extra \"tool\"\n", + "JOIN_DESCRIPTION = (\n", + " \"join():\\n\"\n", + " \" - Collects and combines results from prior actions.\\n\"\n", + " \" - A LLM agent is called upon invoking join to either finalize the user query or wait until the plans are executed.\\n\"\n", + " \" - join should always be the last action in the plan, and will be called in two scenarios:\\n\"\n", + " \" (a) if the answer can be determined by gathering the outputs from tasks to generate the final response.\\n\"\n", + " \" (b) if the answer cannot be determined in the planning phase before you execute the plans. \"\n", + ")\n", + "\n", + "planner_prompt_tmpl_str = (\n", + " \"Given a user query, create a plan to solve it with the utmost parallelizability. \"\n", + " \"Each plan should comprise an action from the following {num_tools} types:\\n\"\n", + " \"{tool_descriptions}\"\n", + " f\"\\n{{num_toolsp1}}. {JOIN_DESCRIPTION}\"\n", + " \"Guidelines:\\n\"\n", + " \" - Each action described above contains input/output types and description.\\n\"\n", + " \" - You must strictly adhere to the input and output types for each action.\\n\"\n", + " \" - The action descriptions contain the guidelines. You MUST strictly follow those guidelines when you use the actions.\\n\"\n", + " \" - Each action in the plan should strictly be one of the above types. Follow the Python conventions for each action.\\n\"\n", + " \" - Each action MUST have a unique ID, which is strictly increasing.\\n\"\n", + " \" - Inputs for actions can either be constants or outputs from preceding actions. \"\n", + " \"In the latter case, use the format $id to denote the ID of the previous action whose output will be the input.\\n\"\n", + " f\" - Always call join as the last action in the plan. Say '{END_OF_PLAN}' after you call join\\n\"\n", + " \" - Ensure the plan maximizes parallelizability.\\n\"\n", + " \" - Only use the provided action types. If a query cannot be addressed using these, invoke the join action for the next steps.\\n\"\n", + " \" - Never introduce new actions other than the ones provided.\\n\\n\"\n", + " \"{replan}\"\n", + " \"{examples}\"\n", + ")\n", + "\n", + "\n", + "def _generate_planner_prompt(\n", + " tools: Sequence[BaseTool],\n", + " example_prompt=str,\n", + "):\n", + " tool_descriptions = \"\\n\".join(\n", + " f\"{i+1}. {tool.name}: {tool.description}\" for i, tool in enumerate(tools)\n", + " )\n", + " planner_prompt_template = ChatPromptTemplate.from_messages(\n", + " [(\"system\", planner_prompt_tmpl_str), (\"user\", \"Question: {input}{context}\")]\n", + " ).partial(\n", + " tool_descriptions=tool_descriptions,\n", + " examples=\"Here are some examples:\\n\\n\" + example_prompt\n", + " if example_prompt\n", + " else \"\",\n", + " num_tools=len(tools),\n", + " num_toolsp1=len(tools) + 1,\n", + " )\n", + "\n", + " return planner_prompt_template\n", + "\n", + "\n", + "def create_planner(\n", + " llm: BaseChatModel,\n", + " example_prompt: str,\n", + " tools: Sequence[BaseTool],\n", + " stop: Optional[list[str]] = None,\n", + "):\n", + " og_planner_prompt = _generate_planner_prompt(tools, example_prompt).partial(\n", + " replan=\"\",\n", + " context=\"\",\n", + " )\n", + " replanner_prompt = _generate_planner_prompt(tools, example_prompt).partial(\n", + " replan=' - You are given \"Previous Plan\" which is the plan that the previous agent created along with the execution results '\n", + " \"(given as Observation) of each plan and a general thought (given as Thought) about the executed results.\"\n", + " 'You MUST use these information to create the next plan under \"Current Plan\".\\n'\n", + " ' - When starting the Current Plan, you should start with \"Thought\" that outlines the strategy for the next plan.\\n'\n", + " \" - In the Current Plan, you should NEVER repeat the actions that are already executed in the Previous Plan.\\n\"\n", + " \" - You must continue the task index from the end of the previous one. Do not repeat task indices.\"\n", + " )\n", + " bound_llm = llm.bind(stop=stop)\n", + " return (\n", + " RunnableBranch(\n", + " ((lambda x: x.get(\"context\") is not None), replanner_prompt),\n", + " og_planner_prompt,\n", + " )\n", + " | bound_llm\n", + " | LLMCompilerPlanParser(tools=tools)\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "7feb5c82-b1a9-40ae-863a-0362fe3ce5ea", + "metadata": {}, + "source": [ + "#### Example usage\n", + "\n", + "Here's an example usage of the planner module." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "3fe074ea-7314-47a7-9a9f-a8e6191ea1f3", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Optional\n", + "\n", + "from langchain.tools import tool\n", + "from langchain_openai import ChatOpenAI\n", + "\n", + "\n", + "@tool\n", + "def get_user_id(first_name: str, last_name: str) -> Optional[int]:\n", + " \"\"\"Query the user IDs of everyone with the provided name.\"\"\"\n", + " student_ids = {\n", + " (\"Eric\", \"Zhang\"): 1432,\n", + " (\"Sam\", \"Van Damm\"): 8523,\n", + " (\"Will\", \"Van Damm\"): 2341,\n", + " }\n", + " return student_ids.get((first_name, last_name))\n", + "\n", + "\n", + "@tool\n", + "def get_scores(class_name: str, user_id: int) -> Optional[str]:\n", + " \"\"\"Query the class registry for grades of the provided user ID.\"\"\"\n", + " return {\n", + " (\"Geology\", 1432): \"A+\",\n", + " (\"Geology\", 8523): \"A\",\n", + " (\"Geology\", 2341): \"B\",\n", + " }.get((class_name, user_id))\n", + "\n", + "\n", + "examples = (\n", + " \"Question: What's the user ID for Johnny Drop Tables?\\n\"\n", + " '1. get_user_id(first_name=\"Johnny\", \"ast_name=\"Drop Tables\")\\n'\n", + " f\"2. join(){END_OF_PLAN}\\n\"\n", + " \"###\\n\"\n", + " \"\\n\"\n", + " \"Question: What was Eric Zhang's score in Calc?\\n\"\n", + " '1. get_user_id(\"Eric\")\\n'\n", + " '2. get_scores(\"calc\", \"$1\")\\n'\n", + " f\"3. join(){END_OF_PLAN}\\n\"\n", + " \"###\\n\"\n", + " \"\\n\"\n", + ")\n", + "\n", + "planner = create_planner(\n", + " ChatOpenAI(model=\"gpt-3.5-turbo\"),\n", + " example_prompt=examples,\n", + " tools=[get_user_id, get_scores],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "55fb0f99-4e59-4e5a-b687-a90bb4d06d39", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{1: {'idx': 1,\n", + " 'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: str) -> Optional[int] - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", + " 'args': {'first_name': 'Sam', 'last_name': 'Van Damm'},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 2: {'idx': 2,\n", + " 'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: str) -> Optional[int] - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", + " 'args': {'first_name': 'Will', 'last_name': 'Van Damm'},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 3: {'idx': 3,\n", + " 'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", + " 'args': {'class_name': 'Calc BC', 'user_id': '$1'},\n", + " 'dependencies': [1],\n", + " 'thought': None},\n", + " 4: {'idx': 4,\n", + " 'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", + " 'args': {'class_name': 'Calc BC', 'user_id': '$2'},\n", + " 'dependencies': [2],\n", + " 'thought': None},\n", + " 5: {'idx': 5,\n", + " 'tool': 'join',\n", + " 'args': (),\n", + " 'dependencies': [1, 2, 3, 4],\n", + " 'thought': None}}" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "tasks = planner.invoke(\n", + " {\"input\": \"What are the Calc BC grades for Sam and Will Van Damm?\"}\n", + ")\n", + "tasks" + ] + }, + { + "cell_type": "markdown", + "id": "5d0e795f-61ff-4553-9823-23e7624ca180", + "metadata": {}, + "source": [ + "## 2. Task Fetching Unit\n", + "\n", + "This component schedules the tasks. In the paper, it's kept separate from the \"executor\", but here we create a single DAG defined in LangChain expression language.\n", + "\n", + "The basic idea is that, given a list of dicts of the form:\n", + "\n", + "```typescript\n", + "{\n", + " tool: BaseTool,\n", + " dependencies: number[],\n", + "}\n", + "```\n", + "\n", + "1. Create a topological sort of the tasks\n", + "2. Execute them on the previous step's output, ensuring to perform variable substitution where appropriate\n", + "\n", + "If we make the assumption that the tasks generated by the LLM are already sorted, we could adapt this to execute in a purely streaming fashion." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "c70a0e28-43db-4ea2-af48-8a2f310dce83", + "metadata": {}, + "outputs": [], + "source": [ + "import functools\n", + "from typing import Any, Union\n", + "\n", + "from langchain_core.runnables import (\n", + " RunnableLambda,\n", + " RunnableParallel,\n", + " RunnablePassthrough,\n", + ")\n", + "\n", + "\n", + "def _sort_tasks(data):\n", + " if not data:\n", + " return []\n", + " sorted_tasks = []\n", + " # Remove tasks already completed\n", + " min_idx = min([int(k) for k in data])\n", + " data = {\n", + " int(k): {\n", + " **v,\n", + " \"dependencies\": [dep for dep in v[\"dependencies\"] if dep >= min_idx],\n", + " }\n", + " for k, v in data.items()\n", + " }\n", + " while data:\n", + " no_deps = {k: v for k, v in data.items() if not v[\"dependencies\"]}\n", + " if not no_deps:\n", + " raise ValueError(\"We seem to have run into a circular dependency.\")\n", + "\n", + " sorted_tasks.append(no_deps)\n", + " data = {\n", + " k: {\n", + " **v,\n", + " \"dependencies\": [d for d in v[\"dependencies\"] if d not in no_deps],\n", + " }\n", + " for k, v in data.items()\n", + " if k not in no_deps\n", + " }\n", + " return sorted_tasks\n", + "\n", + "\n", + "def _resolve_arg(x: dict, arg: Union[str, Any]):\n", + " if isinstance(arg, str) and arg.startswith(\"$\"):\n", + " try:\n", + " return x[f\"task_{arg[1:]}\"]\n", + " except:\n", + " if arg.endswith(\".output\"):\n", + " return x[f\"task_{arg[1:-7]}\"]\n", + " raise\n", + " else:\n", + " return arg\n", + "\n", + "\n", + "def _execute_task(x, task):\n", + " tool_to_use = task[\"tool\"]\n", + " args = task[\"args\"]\n", + " if isinstance(args, str):\n", + " resolved_args = _resolve_arg(x, args)\n", + " elif isinstance(args, dict):\n", + " resolved_args = {key: _resolve_arg(x, val) for key, val in args.items()}\n", + " else:\n", + " # This will likely fail\n", + " resolved_args = args\n", + " try:\n", + " return tool_to_use.invoke(resolved_args)\n", + " except Exception as e:\n", + " return (\n", + " f\"ERROR(Failed to call tool {tool_to_use} with args {tool_to_use}.\"\n", + " + f\" Args resolved to {resolved_args}. Error: {repr(e)})\"\n", + " )\n", + "\n", + "\n", + "def construct_dag(tasks):\n", + " sorted_tasks = _sort_tasks(tasks)\n", + " chain = None\n", + " for idx, task_group in enumerate(sorted_tasks):\n", + " if len(task_group) == 1 and next(iter(task_group.values()))[\"tool\"] == \"join\":\n", + " step = lambda x: {\"join\": x}\n", + " else:\n", + " # Cascade all results forward\n", + " constructor = (\n", + " RunnableParallel if chain is None else RunnablePassthrough.assign\n", + " )\n", + " task_dict = {}\n", + " for idx, task in task_group.items():\n", + " task_dict[f\"task_{idx}\"] = RunnableLambda(\n", + " functools.partial(_execute_task, task=task)\n", + " ).with_config(run_name=f\"task_{idx}\")\n", + "\n", + " step = constructor(**task_dict).with_config(run_name=f\"TaskGroup{idx}\")\n", + " if chain is None:\n", + " chain = step\n", + " else:\n", + " chain |= step\n", + "\n", + " if chain is not None:\n", + " return chain | RunnablePassthrough.assign(tasks=lambda _: tasks)\n", + " return chain" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "87f11cea-3a8d-479c-8a9e-81223dbbc1f5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " +------------------------------+ \n", + " | ParallelInput | \n", + " +------------------------------+ \n", + " *** *** \n", + " ** ** \n", + " ** ** \n", + " +-------------+ +-------------+ \n", + " | Lambda(...) | | Lambda(...) | \n", + " +-------------+ +-------------+ \n", + " *** *** \n", + " ** ** \n", + " ** ** \n", + " +-------------------------------+ \n", + " | ParallelOutput | \n", + " +-------------------------------+ \n", + " * \n", + " * \n", + " * \n", + " +------------------------------+ \n", + " | ParallelInput | \n", + " +------------------------------+ \n", + " ***** * ***** \n", + " ***** * ***** \n", + " *** * *** \n", + "+-------------+ +-------------+ +-------------+ \n", + "| Lambda(...) | | Lambda(...) | | Passthrough | \n", + "+-------------+***** +-------------+ *****+-------------+ \n", + " ***** * ***** \n", + " ***** * ***** \n", + " *** * *** \n", + " +-------------------------------+ \n", + " | ParallelOutput | \n", + " +-------------------------------+ \n", + " * \n", + " * \n", + " * \n", + " +-------------------------------+ \n", + " | Lambda(lambda x: {'join': x}) | \n", + " +-------------------------------+ \n", + " * \n", + " * \n", + " * \n", + " +----------------------+ \n", + " | ParallelInput | \n", + " +----------------------+ \n", + " *** *** \n", + " *** *** \n", + " ** ** \n", + " +-------------------------+ +-------------+ \n", + " | Lambda(lambda _: tasks) | | Passthrough | \n", + " +-------------------------+ +-------------+ \n", + " *** *** \n", + " *** *** \n", + " ** ** \n", + " +-----------------------+ \n", + " | ParallelOutput | \n", + " +-----------------------+ \n" + ] + } + ], + "source": [ + "graph = construct_dag(tasks)\n", + "graph.get_graph().print_ascii()" + ] + }, + { + "cell_type": "markdown", + "id": "9efa15ae-817a-48c6-86ed-16bc112fedc5", + "metadata": {}, + "source": [ + "#### Example Plan\n", + "\n", + "We still haven't introduced any cycles in our computation graph, so this is all easily expressed in LCEL." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "133d124a-0d41-4c6d-a86a-34fa4cb1430f", + "metadata": {}, + "outputs": [], + "source": [ + "chain = planner | construct_dag" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "55142257-2674-4a47-988e-0d2810917329", + "metadata": {}, + "outputs": [], + "source": [ + "example_question = \"Did Sam Van Damm score higher than Eric Zhang in Geology?\"\n", + "task_results = chain.invoke({\"input\": example_question})\n", + "# task_results[\"join\"]" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "9996d98f-0e73-471f-baa8-d8b72d10c7cf", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'join': {'task_1': 8523,\n", + " 'task_2': 1432,\n", + " 'task_3': \"ERROR(Failed to call tool name='get_scores' description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.' args_schema= func= with args name='get_scores' description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.' args_schema= func=. Args resolved to {}. Error: ValidationError(model='get_scoresSchemaSchema', errors=[{'loc': ('class_name',), 'msg': 'field required', 'type': 'value_error.missing'}, {'loc': ('user_id',), 'msg': 'field required', 'type': 'value_error.missing'}]))\",\n", + " 'task_4': \"ERROR(Failed to call tool name='get_scores' description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.' args_schema= func= with args name='get_scores' description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.' args_schema= func=. Args resolved to {}. Error: ValidationError(model='get_scoresSchemaSchema', errors=[{'loc': ('class_name',), 'msg': 'field required', 'type': 'value_error.missing'}, {'loc': ('user_id',), 'msg': 'field required', 'type': 'value_error.missing'}]))\"},\n", + " 'tasks': {1: {'idx': 1,\n", + " 'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: str) -> Optional[int] - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", + " 'args': {'first_name': 'Sam', 'last_name': 'Van Damm'},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 2: {'idx': 2,\n", + " 'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: str) -> Optional[int] - Query the user IDs of everyone with the provided name.', args_schema=, func=),\n", + " 'args': {'first_name': 'Eric', 'last_name': 'Zhang'},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 3: {'idx': 3,\n", + " 'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", + " 'args': {},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 4: {'idx': 4,\n", + " 'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) -> Optional[str] - Query the class registry for grades of the provided user ID.', args_schema=, func=),\n", + " 'args': {},\n", + " 'dependencies': [],\n", + " 'thought': None},\n", + " 5: {'idx': 5,\n", + " 'tool': 'join',\n", + " 'args': (),\n", + " 'dependencies': [1, 2, 3, 4],\n", + " 'thought': None}}}" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "task_results" + ] + }, + { + "cell_type": "markdown", + "id": "563d5311-55f0-4ca1-afbd-01fd970cf3e3", + "metadata": {}, + "source": [ + "## 3. \"Joiner\" \n", + "\n", + "So now we have the planning and initial execution done. We need a component to process these outputs and either:\n", + "\n", + "1. Respond with the correct answer.\n", + "2. Loop with a new plan.\n", + "\n", + "The paper refers to this as the \"joiner\". It's another LLM call, defined below:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "2978991e-45a4-44e6-9deb-f941f44fe93a", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.output_parsers import StrOutputParser\n", + "from typing_extensions import TypedDict\n", + "\n", + "\n", + "def format_task(task, idx):\n", + " tool = task[\"tool\"]\n", + " tool_name = tool if isinstance(tool, str) else tool.name # Handle join()\n", + " args = \", \".join([f\"{k}={v}\" for k, v in task[\"args\"].items()])\n", + " return f\"{idx}. {tool_name}({args})\"\n", + "\n", + "\n", + "def format_tasks(executor_output: dict):\n", + " tasks = executor_output[\"tasks\"]\n", + " prior_observations = executor_output.get(\"observations\")\n", + " joined_output = executor_output[\"join\"]\n", + " execution_results = []\n", + " for idx, task in tasks.items():\n", + " observation_idx = f\"task_{idx}\"\n", + " if observation_idx in joined_output:\n", + " observation = joined_output[observation_idx]\n", + " execution_results.append(f\"{format_task(task, idx)}\\n\\t=> {observation}\")\n", + " joined_results = \"\\n\".join(execution_results)\n", + " result = f\"Executed plan results:\\n{joined_results}\"\n", + " if prior_observations:\n", + " result += f\"\\nPrevious Results:\\n{prior_observations}\"\n", + " return result\n", + "\n", + "\n", + "def _parse_joiner_output(raw_answer: str) -> str:\n", + " thought, answer, is_replan = \"\", \"\", False # default values\n", + " raw_answers = raw_answer.split(\"\\n\")\n", + " for ans in raw_answers:\n", + " if ans.startswith(\"Action:\"):\n", + " answer = ans[ans.find(\"(\") + 1 : ans.find(\")\")]\n", + " is_replan = JOINER_REPLAN in ans\n", + " elif ans.startswith(\"Thought:\"):\n", + " thought = ans.split(\"Thought:\")[1].strip()\n", + " if is_replan:\n", + " return {\"thought\": thought, \"context\": answer}\n", + " else:\n", + " return {\"thought\": thought, \"answer\": answer}" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "b2feea5a-e0e4-4cff-8cb5-fdbfec95ba57", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.prompts import ChatPromptTemplate\n", + "\n", + "\n", + "def create_joiner(prompt, llm):\n", + " return (\n", + " (\n", + " lambda x: {\n", + " **x[\"plan\"],\n", + " \"input\": x[\"input\"],\n", + " \"context\": x.get(\"context\"),\n", + " \"observations\": x.get(\"observations\"),\n", + " }\n", + " )\n", + " | RunnablePassthrough.assign(scratchpad=format_tasks)\n", + " | ChatPromptTemplate.from_messages([(\"system\", prompt), (\"user\", \"{input}\")])\n", + " | llm\n", + " | StrOutputParser()\n", + " | _parse_joiner_output\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "942dab42-ad42-4ba2-90d5-49edbe4fae68", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'thought': 'The initial plan failed to get the scores for Sam Van Damm and Eric Zhang in Geology. In order to provide a final answer, I need these scores.',\n", + " 'context': \"We need to execute get_scores with class_name set to 'Geology' and user_id set to the respective IDs for Sam Van Damm (8523\"}" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "JOINER_FINISH = \"Finish\"\n", + "JOINER_REPLAN = \"Replan\"\n", + "\n", + "system_prompt = (\n", + " \"Solve a question answering task. Here are some guidelines:\\n\"\n", + " \" - In the Assistant Scratchpad, you will be given results of a plan you have executed to answer the user's question.\\n\"\n", + " \" - Thought needs to reason about the question based on the Observations in 1-2 sentences.\\n\"\n", + " \" - Ignore irrelevant action results.\\n\"\n", + " \" - If the required information is present, give a concise but complete and helpful answer to the user's question.\\n\"\n", + " \" - If you are unable to give a satisfactory finishing answer, replan to get the required information.\"\n", + " \" Respond in the following format:\\n\\n\"\n", + " \"Thought: \\n\"\n", + " \"Action: \\n\"\n", + " \"Available actions:\\n\"\n", + " f\" (1) {JOINER_FINISH}(the final answer to return to the user): returns the answer and finishes the task.\\n\"\n", + " f\" (2) {JOINER_REPLAN}(the reasoning and other information that will help you plan again. Can be a line of any length): instructs why we must replan\\n\\n\"\n", + " \" Examples:\\n\"\n", + " \"Question: How many users are currently using the new product?\\n\"\n", + " \"...task returns the number 32,000\\n\"\n", + " \"Thought: I find no issue with the original plan, and the results satisfy everything in the user question.\\n\"\n", + " f\"Action: {JOINER_FINISH}(32,000 users currently use the new product)\\n###\\n\"\n", + " \"Question: How much cooler is it in NY than SF?\\n\"\n", + " \"...task results show SF is 57 degrees fahrenheit today, and they show in NY it has a high of 32 degrees fahrenheit \\n\"\n", + " \"Thought: I can answer by synthesizing the results.\\n\"\n", + " f\"Action: {JOINER_FINISH}(NY is 25 degrees cooler than SF today, as it has a high of 32 degrees Fahrenheit today, whereas in SF, it is 57 degrees Fahrenheit.)\\n###\\n\"\n", + " \"Question: Are the gophers beating the rabbits??\\n\"\n", + " \"...task returns the a score of 7 for rabbits but no other value...\\n\"\n", + " \"Thought: I need the gophers' score to make a final decision.\\n\"\n", + " f\"Action: {JOINER_REPLAN}(The rabbits have a score of 7, but I need the gophers' score.)\"\n", + " \"\\n\\nAssistant Scratchpad:\\n{scratchpad}\"\n", + ")\n", + "\n", + "\n", + "joiner = create_joiner(system_prompt, ChatOpenAI(model=\"gpt-4\"))\n", + "joiner.invoke({\"plan\": task_results, \"input\": example_question})" + ] + }, + { + "cell_type": "markdown", + "id": "b099e5ee-2c23-47d9-9387-0f64e02627d3", + "metadata": {}, + "source": [ + "## Compose using LangGraph\n", + "\n", + "Now we have all the required pieces! Let's construct an LLMCompiler agent. We'll give it a search engine (Tavily) and a simple \"calculate\" function." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "565b08d2-d19c-4125-97f7-996fc01bc631", + "metadata": {}, + "outputs": [], + "source": [ + "import getpass\n", + "import os\n", + "\n", + "os.environ[\"TAVILY_API_KEY\"] = (\n", + " os.environ.get(\"TAVILY_API_KEY\")\n", + " if \"TAVILY_API_KEY\" in os.environ\n", + " else getpass.getpass(\"Tavily API Key:\")\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "bd84ca4b-eacb-471d-9975-5449047b5bed", + "metadata": {}, + "outputs": [], + "source": [ + "from operator import add, mul, sub, truediv\n", + "from typing import Literal\n", + "\n", + "from langchain_community.agent_toolkits import GmailToolkit\n", + "from langchain_community.tools.tavily_search import TavilySearchResults\n", + "from langchain_core.tools import tool\n", + "\n", + "\n", + "@tool\n", + "def calculate(\n", + " arg1: float,\n", + " arg2: float,\n", + " op: Union[Literal[\"+\"], Literal[\"-\"], Literal[\"*\"], Literal[\"/\"]],\n", + "):\n", + " \"\"\"Calculate a mathematical operation on two arguments.\"\"\"\n", + " resolved_op = {\"+\": add, \"-\": sub, \"*\": mul, \"/\": truediv}\n", + " return resolved_op[op](arg1, arg2)\n", + "\n", + "\n", + "tools = [TavilySearchResults(max_results=1), calculate]" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "70566dff-d66e-4c8c-aed3-980522d175c2", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "4.0" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "calculate.invoke(dict(arg1=1, arg2=3, op=\"+\"))" + ] + }, + { + "cell_type": "markdown", + "id": "0ad5d62b-5898-4b9f-808f-c03620c2fe7a", + "metadata": {}, + "source": [ + "#### Defining the stateful graph\n", + "\n", + "We'll define the agent as a stateful graph, with the main nodes being:\n", + "\n", + "1. Plan and execute (the DAG from the first step above)\n", + "2. Join: determine if we should finish or replan\n", + "3. Recontextualize: update the graph state based on the output from the joiner\n" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "c5869ec6-b07f-4fb0-92a7-e6521f0c0bdd", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Dict\n", + "\n", + "MAX_ITERATIONS = 5\n", + "\n", + "\n", + "class GraphState(TypedDict):\n", + " input: str\n", + " plan: Dict\n", + " agent_output: Dict\n", + " observations: Dict\n", + " num_iterations: int # Maximum\n", + " context: str # Exra commentary for the joiner\n", + " stop_reason: str\n", + "\n", + "\n", + "def recontextualize(state):\n", + " # Insert a context string for the re-planner.\n", + " # This could alternatively call an LLM to provide additional logic\n", + " context = state[\"agent_output\"][\"context\"]\n", + " num_iterations = int(state.get(\"num_iterations\") or 1) + 1\n", + " formatted_tasks = format_tasks(state[\"plan\"])\n", + " context_str = f\"\\n\\nPrevious Plan:\\n{formatted_tasks}\\n\" f\"{context}\"\n", + " observations = state[\"observations\"] or {}\n", + " for task, observation in state[\"plan\"][\"join\"].items():\n", + " observations[task] = observation\n", + " return {\n", + " \"context\": context_str,\n", + " \"num_iterations\": num_iterations,\n", + " \"observations\": observations,\n", + " }\n", + "\n", + "\n", + "def add_stop_reason(state: GraphState):\n", + " # Helpful for letting the user know why the agent responded the way it did\n", + " num_iterations = int(state.get(\"num_iterations\") or 0)\n", + " if num_iterations >= MAX_ITERATIONS:\n", + " return {\"stop_reason\": \"end_max_iter\"}\n", + " if state[\"agent_output\"].get(\"answer\"):\n", + " return {\"stop_reason\": \"answer\"}\n", + " return {\"stop_reason\": None}" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "768b5f11-e3d2-47be-8143-a7dcd8765243", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "from langchain_core.tools import BaseTool\n", + "from langgraph.graph import END, StateGraph\n", + "\n", + "workflow = StateGraph(GraphState)\n", + "\n", + "# 1. Define vertices\n", + "\n", + "planner = create_planner(\n", + " llm=ChatOpenAI(model=\"gpt-4-1106-preview\"),\n", + " # Add more examples to improve reliability\n", + " example_prompt=(\n", + " \"Question: What's the capital of Myanmar?\\n\"\n", + " '1. tavily_search_results_json(query=\"Capital of Myanmar)\\n'\n", + " f\"2. join(){END_OF_PLAN}\\n\"\n", + " \"###\\n\"\n", + " \"\\n\"\n", + " ),\n", + " tools=tools,\n", + ")\n", + "\n", + "plan_and_execute = planner | construct_dag\n", + "joiner = create_joiner(system_prompt, ChatOpenAI(model=\"gpt-4-1106-preview\"))\n", + "\n", + "\n", + "# Assign each node to a state variable to update\n", + "workflow.add_node(\"plan_and_execute\", RunnablePassthrough.assign(plan=plan_and_execute))\n", + "workflow.add_node(\"join\", RunnablePassthrough.assign(agent_output=joiner))\n", + "workflow.add_node(\"recontextualize\", recontextualize)\n", + "workflow.add_node(\"provide_stop_reason\", add_stop_reason)\n", + "\n", + "\n", + "## Define edges\n", + "\n", + "workflow.add_edge(\"plan_and_execute\", \"join\")\n", + "workflow.add_edge(\"recontextualize\", \"plan_and_execute\")\n", + "workflow.add_edge(\"join\", \"provide_stop_reason\")\n", + "\n", + "### This condition determines looping logic\n", + "\n", + "\n", + "def should_continue(state):\n", + " if state[\"stop_reason\"] is None:\n", + " return \"continue\"\n", + " return \"end\"\n", + "\n", + "\n", + "workflow.add_conditional_edges(\n", + " start_key=\"provide_stop_reason\",\n", + " # Next, we pass in the function that will determine which node is called next.\n", + " condition=should_continue,\n", + " conditional_edge_mapping={\n", + " # If it generates context, we must replan\n", + " \"continue\": \"recontextualize\",\n", + " # Otherwise we finish.\n", + " \"end\": END,\n", + " },\n", + ")\n", + "workflow.set_entry_point(\"plan_and_execute\")\n", + "chain = workflow.compile()" + ] + }, + { + "cell_type": "markdown", + "id": "9f8c9849-8531-463d-a0ef-dcc3d9888b2d", + "metadata": {}, + "source": [ + "## Simple question\n", + "\n", + "Let's ask a simple question of the agent." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "5bc4584a-e31c-4065-805e-76a6db30676a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "In 2022, the real GDP of New York was about 1.56 trillion U.S. dollars.\n" + ] + } + ], + "source": [ + "result = chain.invoke({\"input\": \"What's the GDP of New York?\"})\n", + "print(result[\"agent_output\"][\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "id": "33c65ef5-b4b2-4ab2-8c78-a551da7819b9", + "metadata": {}, + "source": [ + "## Multi-hop question" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "0b3a0916-d8ca-4092-b91c-d9e2b05259d8", + "metadata": {}, + "outputs": [], + "source": [ + "result = chain.invoke(\n", + " {\n", + " \"input\": \"What's the oldest parrot alive, and how much longer is that than the average?\"\n", + " },\n", + " {\n", + " \"recursion_limit\": 100,\n", + " },\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "6c65c414-7668-4fdf-ba97-f42f659b1317", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Cookie, a cockatoo, was the oldest parrot alive, having reached the age of 83, which is 23 years longer than the maximum average lifespan of a cockatoo in captivity, which is 60 years.\n" + ] + } + ], + "source": [ + "print(result[\"agent_output\"][\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "id": "1b859bc7-1a85-4d35-b57b-f67c87282403", + "metadata": {}, + "source": [ + "## Streaming" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "38d3ea91-59ba-4267-8060-ed75bbc840c6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Step: {'plan_and...\n", + "Step: {'join': {...\n", + "Step: {'provide_...\n", + "Step: {'recontex...\n", + "Step: {'plan_and...\n", + "Step: {'join': {...\n", + "Step: {'provide_...\n", + "Step: {'__end__'...\n", + "3307.0\n" + ] + } + ], + "source": [ + "last_step = None\n", + "for step in chain.stream({\"input\": \"What's ((3*(4+5)/0.5)+3245) + 8?\"}):\n", + " print(\"Step: \", str(step)[:10] + \"...\")\n", + " last_step = step\n", + "print(\"***\")\n", + "print(last_step[\"__end__\"][\"agent_output\"][\"answer\"])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/advanced_agents/llm_compiler/output_parser.py b/examples/advanced_agents/llm_compiler/output_parser.py new file mode 100644 index 000000000..e8590bb17 --- /dev/null +++ b/examples/advanced_agents/llm_compiler/output_parser.py @@ -0,0 +1,154 @@ +import re +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union +import ast + +from langchain_core.exceptions import OutputParserException +from langchain_core.messages import BaseMessage +from langchain_core.output_parsers import BaseOutputParser +from langchain_core.tools import BaseTool + +THOUGHT_PATTERN = r"Thought: ([^\n]*)" +ACTION_PATTERN = r"\n*(\d+)\. (\w+)\((.*)\)(\s*#\w+\n)?" +# $1 or ${1} -> 1 +ID_PATTERN = r"\$\{?(\d+)\}?" +END_OF_PLAN = "" + + +### Helper functions + + +def _ast_parse(arg: str) -> Any: + try: + return ast.literal_eval(arg) + except: # noqa + return arg + + +def _parse_llm_compiler_action_args(args: str, tool: Union[str, BaseTool]) -> list[Any]: + """Parse arguments from a string.""" + if args == "": + return () + if isinstance(tool, str): + return () + extracted_args = {} + tool_key = None + prev_idx = None + for key in tool.args.keys(): + # Split if present + if f"{key}=" in args: + idx = args.index(f"{key}=") + if prev_idx is not None: + # print(tool_key, flush=True) + # breakpoint() + extracted_args[tool_key] = _ast_parse( + args[prev_idx:idx].strip().rstrip(",") + ) + args = args.split(f"{key}=", 1)[1] + tool_key = key + prev_idx = 0 + if prev_idx is not None: + extracted_args[tool_key] = _ast_parse( + args[prev_idx:].strip().rstrip(",").rstrip(")") + ) + return extracted_args + + +def default_dependency_rule(idx, args: str): + matches = re.findall(ID_PATTERN, args) + numbers = [int(match) for match in matches] + return idx in numbers + + +def _get_dependencies_from_graph( + idx: int, tool_name: str, args: Dict[str, Any] +) -> dict[str, list[str]]: + """Get dependencies from a graph.""" + if tool_name == "join": + return list(range(1, idx)) + return [i for i in range(1, idx) if default_dependency_rule(i, str(args))] + + +def instantiate_task( + tools: Sequence[BaseTool], + idx: int, + tool_name: str, + args: Union[str, Any], + thought: Optional[str] = None, +) -> dict: + if tool_name == "join": + tool = "join" + else: + try: + tool = tools[[tool.name for tool in tools].index(tool_name)] + except ValueError as e: + raise OutputParserException(f"Tool {tool_name} not found.") from e + tool_args = _parse_llm_compiler_action_args(args, tool) + dependencies = _get_dependencies_from_graph(idx, tool_name, tool_args) + + return dict( + idx=idx, + tool=tool, + args=tool_args, + dependencies=dependencies, + thought=thought, + ) + + +class LLMCompilerPlanParser(BaseOutputParser[dict], extra="allow"): + """Planning output parser.""" + + tools: List[BaseTool] + + def _transform(self, input: Iterator[Union[str, BaseMessage]]) -> Iterator[dict]: + texts = [] + thought = None + for chunk in input: + # Assume input is str. TODO: support vision/other formats + text = chunk if isinstance(chunk, str) else str(chunk.content) + for task, thought in self.ingest_token(text, texts, thought): + yield {task["idx"]: task} + # Final possible task + if texts: + task, _ = self._parse_task("".join(texts), thought) + if task: + yield {task["idx"]: task} + + def parse(self, text: str): + task_dict = {} + for task in self._transform([text]): + task_dict.update(task) + return task_dict + + def ingest_token( + self, token: str, buffer: List[str], thought: Optional[str] + ) -> Iterator[Tuple[Optional[dict], str]]: + buffer.append(token) + if "\n" in token: + buffer_ = "".join(buffer).split("\n") + suffix = buffer_[-1] + for line in buffer_[:-1]: + task, thought = self._parse_task(line, thought) + if task: + yield task, thought + buffer.clear() + buffer.append(suffix) + + def _parse_task(self, line: str, thought: Optional[str] = None): + task = None + if match := re.match(THOUGHT_PATTERN, line): + # Optionally, action can be preceded by a thought + thought = match.group(1) + elif match := re.match(ACTION_PATTERN, line): + # if action is parsed, return the task, and clear the buffer + idx, tool_name, args, _ = match.groups() + idx = int(idx) + task = instantiate_task( + tools=self.tools, + idx=idx, + tool_name=tool_name, + args=args, + thought=thought, + ) + thought = None + # Else it is just dropped + return task, thought