examples/llm-compiler/LLMCompiler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c8b472b-f3fb-46c2-841f-930a4692697b",
   "metadata": {},
   "source": [
    "# LLMCompiler\n",
    "\n",
    "This notebook shows how to implement [LLMCompiler, by Kim, et. al](https://arxiv.org/abs/2312.04511) in LangGraph.\n",
    "\n",
    "LLMCompiler is an agent architecture designed to **speed up** the execution of agentic tasks by eagerly-executed tasks within a DAG. It also saves costs on redundant token usage by reducing the number of calls to the LLM. Below is an overview of its computational graph:\n",
    "\n",
    "![LLMCompiler Graph](./img/llm-compiler.png)\n",
    "\n",
    "It has 3 main components:\n",
    "\n",
    "1. Planner: stream a DAG of tasks.\n",
    "2. Task Fetching Unit: schedules and executes the tasks as soon as they are executable\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. The end result will leave a trace [like the following](https://smith.langchain.com/public/218c2677-c719-4147-b0e9-7bc3b5bb2623/r).\n",
    "\n",
    "\n",
    "**First,** install the dependencies, and set up LangSmith for tracing to more easily debug and observe the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "16bd5497-35ad-44f2-94d9-19ff39a5ffed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install -U --quiet langchain_openai langsmith langgraph langchain numexpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "abbd6948-e9a3-47ca-89c7-7ac2fc5eca8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "\n",
    "def _get_pass(var: str):\n",
    "    if var not in os.environ:\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "# Optional: Debug + trace calls using LangSmith\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"True\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"LLMCompiler\"\n",
    "_get_pass(\"LANGCHAIN_API_KEY\")\n",
    "_get_pass(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a61b48ee-8c6f-4863-913a-676f659287de",
   "metadata": {},
   "source": [
    "## Part 1: Tools\n",
    "\n",
    "We'll first define the tools for the agent to use in our demo. We'll give it the class search engine + calculator combo.\n",
    "\n",
    "If you don't want to sign up for tavily, you can replace it with the free [DuckDuckGo](https://python.langchain.com/docs/integrations/tools/ddg)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e7476bb2-1a51-42f6-b7ae-82a0300bbf84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "\n",
    "# Imported from the https://github.com/langchain-ai/langgraph/tree/main/examples/plan-and-execute repo\n",
    "from math_tools import get_math_tool\n",
    "\n",
    "_get_pass(\"TAVILY_API_KEY\")\n",
    "\n",
    "calculate = get_math_tool(ChatOpenAI(model=\"gpt-4-turbo-preview\"))\n",
    "search = TavilySearchResults(\n",
    "    max_results=1,\n",
    "    description='tavily_search_results_json(query=\"the search query\") - a search engine.',\n",
    ")\n",
    "\n",
    "tools = [search, calculate]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "152eecf3-6bef-4718-af71-a0b3c5a3b009",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'37'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate.invoke(\n",
    "    {\n",
    "        \"problem\": \"What's the temp of sf + 5?\",\n",
    "        \"context\": [\"Thet empreature of sf is 32 degrees\"],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1abdedbd-d81b-4ee9-b46f-f29439ed1350",
   "metadata": {},
   "source": [
    "# Part 2: 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()<END_OF_PLAN>\"\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": 5,
   "id": "15dd9639-691f-4906-9012-83fd6e9ac126",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "Given a user query, create a plan to solve it with the utmost parallelizability. Each plan should comprise an action from the following \u001b[33;1m\u001b[1;3m{num_tools}\u001b[0m types:\n",
      "\u001b[33;1m\u001b[1;3m{tool_descriptions}\u001b[0m\n",
      "\u001b[33;1m\u001b[1;3m{num_tools}\u001b[0m. join(): Collects and combines results from prior actions.\n",
      "\n",
      " - An LLM agent is called upon invoking join() to either finalize the user query or wait until the plans are executed.\n",
      " - join should always be the last action in the plan, and will be called in two scenarios:\n",
      "   (a) if the answer can be determined by gathering the outputs from tasks to generate the final response.\n",
      "   (b) if the answer cannot be determined in the planning phase before you execute the plans. Guidelines:\n",
      " - Each action described above contains input/output types and description.\n",
      "    - You must strictly adhere to the input and output types for each action.\n",
      "    - The action descriptions contain the guidelines. You MUST strictly follow those guidelines when you use the actions.\n",
      " - Each action in the plan should strictly be one of the above types. Follow the Python conventions for each action.\n",
      " - Each action MUST have a unique ID, which is strictly increasing.\n",
      " - Inputs for actions can either be constants or outputs from preceding actions. In the latter case, use the format $id to denote the ID of the previous action whose output will be the input.\n",
      " - Always call join as the last action in the plan. Say '<END_OF_PLAN>' after you call join\n",
      " - Ensure the plan maximizes parallelizability.\n",
      " - Only use the provided action types. If a query cannot be addressed using these, invoke the join action for the next steps.\n",
      " - Never introduce new actions other than the ones provided.\n",
      "\n",
      "=============================\u001b[1m Messages Placeholder \u001b[0m=============================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{messages}\u001b[0m\n",
      "\n",
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "Remember, ONLY respond with the task list in the correct format! E.g.:\n",
      "idx. tool(arg_name=args)\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "from typing import Sequence\n",
    "\n",
    "from langchain_core.language_models 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 langchain_core.messages import (\n",
    "    BaseMessage,\n",
    "    FunctionMessage,\n",
    "    HumanMessage,\n",
    "    SystemMessage,\n",
    ")\n",
    "\n",
    "from output_parser import LLMCompilerPlanParser, Task\n",
    "from langchain import hub\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "prompt = hub.pull(\"wfh/llm-compiler\")\n",
    "print(prompt.pretty_print())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "45689d40-d8df-4316-a121-6ea9c87d2efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_planner(\n",
    "    llm: BaseChatModel, tools: Sequence[BaseTool], base_prompt: ChatPromptTemplate\n",
    "):\n",
    "    tool_descriptions = \"\\n\".join(\n",
    "        f\"{i+1}. {tool.description}\\n\"\n",
    "        for i, tool in enumerate(\n",
    "            tools\n",
    "        )  # +1 to offset the 0 starting index, we want it count normally from 1.\n",
    "    )\n",
    "    planner_prompt = base_prompt.partial(\n",
    "        replan=\"\",\n",
    "        num_tools=len(tools)\n",
    "        + 1,  # Add one because we're adding the join() tool at the end.\n",
    "        tool_descriptions=tool_descriptions,\n",
    "    )\n",
    "    replanner_prompt = base_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",
    "        num_tools=len(tools) + 1,\n",
    "        tool_descriptions=tool_descriptions,\n",
    "    )\n",
    "\n",
    "    def should_replan(state: list):\n",
    "        # Context is passed as a system message\n",
    "        return isinstance(state[-1], SystemMessage)\n",
    "\n",
    "    def wrap_messages(state: list):\n",
    "        return {\"messages\": state}\n",
    "\n",
    "    def wrap_and_get_last_index(state: list):\n",
    "        next_task = 0\n",
    "        for message in state[::-1]:\n",
    "            if isinstance(message, FunctionMessage):\n",
    "                next_task = message.additional_kwargs[\"idx\"] + 1\n",
    "                break\n",
    "        state[-1].content = state[-1].content + f\" - Begin counting at : {next_task}\"\n",
    "        return {\"messages\": state}\n",
    "\n",
    "    return (\n",
    "        RunnableBranch(\n",
    "            (should_replan, wrap_and_get_last_index | replanner_prompt),\n",
    "            wrap_messages | planner_prompt,\n",
    "        )\n",
    "        | llm\n",
    "        | LLMCompilerPlanParser(tools=tools)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bbdcb57b-5362-4b9e-88db-fb3fae443fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")\n",
    "# This is the primary \"agent\" in our application\n",
    "planner = create_planner(llm, tools, prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "730490c6-6e3a-4173-82a1-9eb9d5eeff20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "description='tavily_search_results_json(query=\"the search query\") - a search engine.' max_results=1 {'query': 'current temperature in San Francisco'}\n",
      "---\n",
      "name='math' description='math(problem: str, context: Optional[List[str]] = None, config: Optional[langchain_core.runnables.config.RunnableConfig] = None) - math(problem: str, context: Optional[list[str]]) -> float:\\n - Solves the provided math problem.\\n - `problem` can be either a simple math problem (e.g. \"1 + 3\") or a word problem (e.g. \"how many apples are there if there are 3 apples and 2 apples\").\\n - You cannot calculate multiple expressions in one call. For instance, `math(\\'1 + 3, 2 + 4\\')` does not work. If you need to calculate multiple expressions, you need to call them separately like `math(\\'1 + 3\\')` and then `math(\\'2 + 4\\')`\\n - Minimize the number of `math` actions as much as possible. For instance, instead of calling 2. math(\"what is the 10% of $1\") and then call 3. math(\"$1 + $2\"), you MUST call 2. math(\"what is the 110% of $1\") instead, which will reduce the number of math actions.\\n - You can optionally provide a list of strings as `context` to help the agent solve the problem. If there are multiple contexts you need to answer the question, you can provide them as a list of strings.\\n - `math` action will not see the output of the previous actions unless you provide it as `context`. You MUST provide the output of the previous actions as `context` if you need to do math on it.\\n - You MUST NEVER provide `search` type action\\'s outputs as a variable in the `problem` argument. This is because `search` returns a text blob that contains the information about the entity, not a number or value. Therefore, when you need to provide an output of `search` action, you MUST provide it as a `context` argument to `math` action. For example, 1. search(\"Barack Obama\") and then 2. math(\"age of $1\") is NEVER allowed. Use 2. math(\"age of Barack Obama\", context=[\"$1\"]) instead.\\n - When you ask a question about `context`, specify the units. For instance, \"what is xx in height?\" or \"what is xx in millions?\" instead of \"what is xx?\"' args_schema=<class 'pydantic.v1.main.mathSchema'> func=<function get_math_tool.<locals>.calculate_expression at 0x10f354ea0> {'problem': 'raise $0 to the 3rd power', 'context': ['$0']}\n",
      "---\n",
      "join ()\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "example_question = \"What's the temperature in SF raised to the 3rd power?\"\n",
    "\n",
    "for task in planner.stream([HumanMessage(content=example_question)]):\n",
    "    print(task[\"tool\"], task[\"args\"])\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d0e795f-61ff-4553-9823-23e7624ca180",
   "metadata": {},
   "source": [
    "## 3. Task Fetching Unit\n",
    "\n",
    "This component schedules the tasks. It receives a stream of tools of the following format:\n",
    "\n",
    "```typescript\n",
    "{\n",
    "    tool: BaseTool,\n",
    "    dependencies: number[],\n",
    "}\n",
    "```\n",
    "\n",
    "\n",
    "The basic idea is to begin executing tools as soon as their dependencies are met. This is done through multi-threading. We will combine the task fetching unit and executor below:\n",
    "\n",
    "![diagram](./img/diagram.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c1fbafdd-42d4-4575-8466-e5951cee71f4",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "from typing import Any, Union, Iterable, List, Tuple, Dict\n",
    "from typing_extensions import TypedDict\n",
    "import re\n",
    "\n",
    "from langchain_core.runnables import (\n",
    "    chain as as_runnable,\n",
    ")\n",
    "\n",
    "from concurrent.futures import ThreadPoolExecutor, wait\n",
    "import time\n",
    "\n",
    "\n",
    "def _get_observations(messages: List[BaseMessage]) -> Dict[int, Any]:\n",
    "    # Get all previous tool responses\n",
    "    results = {}\n",
    "    for message in messages[::-1]:\n",
    "        if isinstance(message, FunctionMessage):\n",
    "            results[int(message.additional_kwargs[\"idx\"])] = message.content\n",
    "    return results\n",
    "\n",
    "\n",
    "class SchedulerInput(TypedDict):\n",
    "    messages: List[BaseMessage]\n",
    "    tasks: Iterable[Task]\n",
    "\n",
    "\n",
    "def _execute_task(task, observations, config):\n",
    "    tool_to_use = task[\"tool\"]\n",
    "    if isinstance(tool_to_use, str):\n",
    "        return tool_to_use\n",
    "    args = task[\"args\"]\n",
    "    try:\n",
    "        if isinstance(args, str):\n",
    "            resolved_args = _resolve_arg(args, observations)\n",
    "        elif isinstance(args, dict):\n",
    "            resolved_args = {\n",
    "                key: _resolve_arg(val, observations) for key, val in args.items()\n",
    "            }\n",
    "        else:\n",
    "            # This will likely fail\n",
    "            resolved_args = args\n",
    "    except Exception as e:\n",
    "        return (\n",
    "            f\"ERROR(Failed to call {tool_to_use.name} with args {args}.)\"\n",
    "            f\" Args could not be resolved. Error: {repr(e)}\"\n",
    "        )\n",
    "    try:\n",
    "        return tool_to_use.invoke(resolved_args, config)\n",
    "    except Exception as e:\n",
    "        return (\n",
    "            f\"ERROR(Failed to call {tool_to_use.name} with args {args}.\"\n",
    "            + f\" Args resolved to {resolved_args}. Error: {repr(e)})\"\n",
    "        )\n",
    "\n",
    "\n",
    "def _resolve_arg(arg: Union[str, Any], observations: Dict[int, Any]):\n",
    "    # $1 or ${1} -> 1\n",
    "    ID_PATTERN = r\"\\$\\{?(\\d+)\\}?\"\n",
    "\n",
    "    def replace_match(match):\n",
    "        # If the string is ${123}, match.group(0) is ${123}, and match.group(1) is 123.\n",
    "\n",
    "        # Return the match group, in this case the index, from the string. This is the index\n",
    "        # number we get back.\n",
    "        idx = int(match.group(1))\n",
    "        return str(observations.get(idx, match.group(0)))\n",
    "\n",
    "    # For dependencies on other tasks\n",
    "    if isinstance(arg, str):\n",
    "        return re.sub(ID_PATTERN, replace_match, arg)\n",
    "    elif isinstance(arg, list):\n",
    "        return [_resolve_arg(a, observations) for a in arg]\n",
    "    else:\n",
    "        return str(arg)\n",
    "\n",
    "\n",
    "@as_runnable\n",
    "def schedule_task(task_inputs, config):\n",
    "    task: Task = task_inputs[\"task\"]\n",
    "    observations: Dict[int, Any] = task_inputs[\"observations\"]\n",
    "    try:\n",
    "        observation = _execute_task(task, observations, config)\n",
    "    except Exception:\n",
    "        import traceback\n",
    "\n",
    "        observation = traceback.format_exception()  # repr(e) +\n",
    "    observations[task[\"idx\"]] = observation\n",
    "\n",
    "\n",
    "def schedule_pending_task(\n",
    "    task: Task, observations: Dict[int, Any], retry_after: float = 0.2\n",
    "):\n",
    "    while True:\n",
    "        deps = task[\"dependencies\"]\n",
    "        if deps and (any([dep not in observations for dep in deps])):\n",
    "            # Dependencies not yet satisfied\n",
    "            time.sleep(retry_after)\n",
    "            continue\n",
    "        schedule_task.invoke({\"task\": task, \"observations\": observations})\n",
    "        break\n",
    "\n",
    "\n",
    "@as_runnable\n",
    "def schedule_tasks(scheduler_input: SchedulerInput) -> List[FunctionMessage]:\n",
    "    \"\"\"Group the tasks into a DAG schedule.\"\"\"\n",
    "    # For streaming, we are making a few simplifying assumption:\n",
    "    # 1. The LLM does not create cyclic dependencies\n",
    "    # 2. That the LLM will not generate tasks with future deps\n",
    "    # If this ceases to be a good assumption, you can either\n",
    "    # adjust to do a proper topological sort (not-stream)\n",
    "    # or use a more complicated data structure\n",
    "    tasks = scheduler_input[\"tasks\"]\n",
    "    args_for_tasks = {}\n",
    "    messages = scheduler_input[\"messages\"]\n",
    "    # If we are re-planning, we may have calls that depend on previous\n",
    "    # plans. Start with those.\n",
    "    observations = _get_observations(messages)\n",
    "    task_names = {}\n",
    "    originals = set(observations)\n",
    "    # ^^ We assume each task inserts a different key above to\n",
    "    # avoid race conditions...\n",
    "    futures = []\n",
    "    retry_after = 0.25  # Retry every quarter second\n",
    "    with ThreadPoolExecutor() as executor:\n",
    "        for task in tasks:\n",
    "            deps = task[\"dependencies\"]\n",
    "            task_names[task[\"idx\"]] = (\n",
    "                task[\"tool\"] if isinstance(task[\"tool\"], str) else task[\"tool\"].name\n",
    "            )\n",
    "            args_for_tasks[task[\"idx\"]] = task[\"args\"]\n",
    "            if (\n",
    "                # Depends on other tasks\n",
    "                deps\n",
    "                and (any([dep not in observations for dep in deps]))\n",
    "            ):\n",
    "                futures.append(\n",
    "                    executor.submit(\n",
    "                        schedule_pending_task, task, observations, retry_after\n",
    "                    )\n",
    "                )\n",
    "            else:\n",
    "                # No deps or all deps satisfied\n",
    "                # can schedule now\n",
    "                schedule_task.invoke(dict(task=task, observations=observations))\n",
    "                # futures.append(executor.submit(schedule_task.invoke dict(task=task, observations=observations)))\n",
    "\n",
    "        # All tasks have been submitted or enqueued\n",
    "        # Wait for them to complete\n",
    "        wait(futures)\n",
    "    # Convert observations to new tool messages to add to the state\n",
    "    new_observations = {\n",
    "        k: (task_names[k], args_for_tasks[k], observations[k])\n",
    "        for k in sorted(observations.keys() - originals)\n",
    "    }\n",
    "    tool_messages = [\n",
    "        FunctionMessage(\n",
    "            name=name, content=str(obs), additional_kwargs={\"idx\": k, \"args\": task_args}\n",
    "        )\n",
    "        for k, (name, task_args, obs) in new_observations.items()\n",
    "    ]\n",
    "    return tool_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "052f6b16-103a-40e9-94dd-8fcc37e77ba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "\n",
    "@as_runnable\n",
    "def plan_and_schedule(messages: List[BaseMessage], config):\n",
    "    tasks = planner.stream(messages, config)\n",
    "    # Begin executing the planner immediately\n",
    "    try:\n",
    "        tasks = itertools.chain([next(tasks)], tasks)\n",
    "    except StopIteration:\n",
    "        # Handle the case where tasks is empty.\n",
    "        tasks = iter([])\n",
    "    scheduled_tasks = schedule_tasks.invoke(\n",
    "        {\n",
    "            \"messages\": messages,\n",
    "            \"tasks\": tasks,\n",
    "        },\n",
    "        config,\n",
    "    )\n",
    "    return scheduled_tasks"
   ]
  },
  {
   "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": 11,
   "id": "55142257-2674-4a47-988e-0d2810917329",
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_messages = plan_and_schedule.invoke([HumanMessage(content=example_question)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a98e0525-2fcf-4fa1-baf6-79858bb8a6bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[FunctionMessage(content='[]', additional_kwargs={'idx': 0}, name='tavily_search_results_json'),\n",
       " FunctionMessage(content='ValueError(\\'Failed to evaluate \"N/A\". Raised error: KeyError(\\\\\\'A\\\\\\'). Please try again with a valid numerical expression\\')', additional_kwargs={'idx': 1}, name='math'),\n",
       " FunctionMessage(content='join', additional_kwargs={'idx': 2}, name='join')]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tool_messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563d5311-55f0-4ca1-afbd-01fd970cf3e3",
   "metadata": {},
   "source": [
    "## 4. \"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. We are using function calling to improve parsing reliability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "942dab42-ad42-4ba2-90d5-49edbe4fae68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain.chains.openai_functions import create_structured_output_runnable\n",
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "\n",
    "class FinalResponse(BaseModel):\n",
    "    \"\"\"The final response/answer.\"\"\"\n",
    "\n",
    "    response: str\n",
    "\n",
    "\n",
    "class Replan(BaseModel):\n",
    "    feedback: str = Field(\n",
    "        description=\"Analysis of the previous attempts and recommendations on what needs to be fixed.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class JoinOutputs(BaseModel):\n",
    "    \"\"\"Decide whether to replan or whether you can return the final response.\"\"\"\n",
    "\n",
    "    thought: str = Field(\n",
    "        description=\"The chain of thought reasoning for the selected action\"\n",
    "    )\n",
    "    action: Union[FinalResponse, Replan]\n",
    "\n",
    "\n",
    "joiner_prompt = hub.pull(\"wfh/llm-compiler-joiner\").partial(\n",
    "    examples=\"\"\n",
    ")  # You can optionally add examples\n",
    "llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")\n",
    "\n",
    "runnable = create_structured_output_runnable(JoinOutputs, llm, joiner_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb50c4cd-947c-4a5d-a9f7-f0d92a10600f",
   "metadata": {},
   "source": [
    "We will select only the most recent messages in the state, and format the output to be more useful for\n",
    "the planner, should the agent need to loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "951a33cf-2a05-4a33-899a-0ab1d97122fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _parse_joiner_output(decision: JoinOutputs) -> List[BaseMessage]:\n",
    "    response = [AIMessage(content=f\"Thought: {decision.thought}\")]\n",
    "    if isinstance(decision.action, Replan):\n",
    "        return response + [\n",
    "            SystemMessage(\n",
    "                content=f\"Context from last attempt: {decision.action.feedback}\"\n",
    "            )\n",
    "        ]\n",
    "    else:\n",
    "        return response + [AIMessage(content=decision.action.response)]\n",
    "\n",
    "\n",
    "def select_recent_messages(messages: list) -> dict:\n",
    "    selected = []\n",
    "    for msg in messages[::-1]:\n",
    "        selected.append(msg)\n",
    "        if isinstance(msg, HumanMessage):\n",
    "            break\n",
    "    return {\"messages\": selected[::-1]}\n",
    "\n",
    "\n",
    "joiner = select_recent_messages | runnable | _parse_joiner_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1e49d4b1-8266-4520-a566-1448b1c31c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_messages = [HumanMessage(content=example_question)] + tool_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "31854dfd-b82f-4c24-9b58-6bae66777909",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AIMessage(content='Thought: The search did not return any results, and the attempt to calculate the temperature in San Francisco raised to the 3rd power failed due to missing temperature information.'),\n",
       " SystemMessage(content='Context from last attempt: I need to find the current temperature in San Francisco before calculating its value raised to the 3rd power.')]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joiner.invoke(input_messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b099e5ee-2c23-47d9-9387-0f64e02627d3",
   "metadata": {},
   "source": [
    "## 5. Compose using LangGraph\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "768b5f11-e3d2-47be-8143-a7dcd8765243",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import MessageGraph, END\n",
    "from typing import Dict\n",
    "\n",
    "graph_builder = MessageGraph()\n",
    "\n",
    "# 1.  Define vertices\n",
    "# We defined plan_and_schedule above already\n",
    "# Assign each node to a state variable to update\n",
    "graph_builder.add_node(\"plan_and_schedule\", plan_and_schedule)\n",
    "graph_builder.add_node(\"join\", joiner)\n",
    "\n",
    "\n",
    "## Define edges\n",
    "graph_builder.add_edge(\"plan_and_schedule\", \"join\")\n",
    "\n",
    "### This condition determines looping logic\n",
    "\n",
    "\n",
    "def should_continue(state: List[BaseMessage]):\n",
    "    if isinstance(state[-1], AIMessage):\n",
    "        return END\n",
    "    return \"plan_and_schedule\"\n",
    "\n",
    "\n",
    "graph_builder.add_conditional_edges(\n",
    "    start_key=\"join\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    condition=should_continue,\n",
    ")\n",
    "graph_builder.set_entry_point(\"plan_and_schedule\")\n",
    "chain = graph_builder.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": [
      "{'plan_and_schedule': [FunctionMessage(content='[{\\'url\\': \\'https://www.governor.ny.gov/programs/fy-2024-new-york-state-budget\\', \\'content\\': \"The $229 billion FY 2024 New York State Budget reflects Governor Hochul\\'s bold agenda to make New York more affordable,  FY 2024 Budget Assets FY 2024 New York State Budget Highlights Improving Public Safety  GOVERNOR HOME GOVERNOR KATHY HOCHUL FY 2024 New York State Budget  Transformative investments to support New York\\'s business community and boost the state economy.The $229 billion FY 2024 NYS Budget reflects Governor Hochul\\'s bold agenda to make New York more affordable, more livable, and safer.\"}]', additional_kwargs={'idx': 0}, name='tavily_search_results_json')]}\n",
      "---\n",
      "{'join': [AIMessage(content=\"Thought: The information provided does not specify the Gross Domestic Product (GDP) of New York, but instead provides details about the state's budget for fiscal year 2024, which is $229 billion. This budget figure cannot be accurately equated to the GDP.\"), SystemMessage(content=\"Context from last attempt: The search results provided information about New York's state budget rather than its GDP. To answer the user's question, we need to find specific data on New York's GDP, not its budget.\")]}\n",
      "---\n",
      "{'plan_and_schedule': [FunctionMessage(content=\"[{'url': 'https://en.wikipedia.org/wiki/Economy_of_New_York_(state)', 'content': 'The economy of the State of New York is reflected in its gross state product in 2022 of $2.053 trillion, ranking third  Contents Economy of New York (state)  New York City-centered metropolitan statistical area produced a gross metropolitan product (GMP) of $US2.0 trillion,  of the items in which New York ranks high nationally:The economy of the State of New York is reflected in its gross state product in 2022 of $2.053 trillion, ranking third in size behind the larger states of\\\\xa0...'}]\", additional_kwargs={'idx': 1}, name='tavily_search_results_json')]}\n",
      "---\n",
      "{'join': [AIMessage(content=\"Thought: The required information about New York's GDP is provided in the search results. In 2022, New York had a Gross State Product (GSP) of $2.053 trillion.\"), AIMessage(content='The Gross Domestic Product (GDP) of New York in 2022 was $2.053 trillion.')]}\n",
      "---\n",
      "{'__end__': [HumanMessage(content=\"What's the GDP of New York?\"), FunctionMessage(content='[{\\'url\\': \\'https://www.governor.ny.gov/programs/fy-2024-new-york-state-budget\\', \\'content\\': \"The $229 billion FY 2024 New York State Budget reflects Governor Hochul\\'s bold agenda to make New York more affordable,  FY 2024 Budget Assets FY 2024 New York State Budget Highlights Improving Public Safety  GOVERNOR HOME GOVERNOR KATHY HOCHUL FY 2024 New York State Budget  Transformative investments to support New York\\'s business community and boost the state economy.The $229 billion FY 2024 NYS Budget reflects Governor Hochul\\'s bold agenda to make New York more affordable, more livable, and safer.\"}]', additional_kwargs={'idx': 0}, name='tavily_search_results_json'), AIMessage(content=\"Thought: The information provided does not specify the Gross Domestic Product (GDP) of New York, but instead provides details about the state's budget for fiscal year 2024, which is $229 billion. This budget figure cannot be accurately equated to the GDP.\"), SystemMessage(content=\"Context from last attempt: The search results provided information about New York's state budget rather than its GDP. To answer the user's question, we need to find specific data on New York's GDP, not its budget. - Begin counting at : 1\"), FunctionMessage(content=\"[{'url': 'https://en.wikipedia.org/wiki/Economy_of_New_York_(state)', 'content': 'The economy of the State of New York is reflected in its gross state product in 2022 of $2.053 trillion, ranking third  Contents Economy of New York (state)  New York City-centered metropolitan statistical area produced a gross metropolitan product (GMP) of $US2.0 trillion,  of the items in which New York ranks high nationally:The economy of the State of New York is reflected in its gross state product in 2022 of $2.053 trillion, ranking third in size behind the larger states of\\\\xa0...'}]\", additional_kwargs={'idx': 1}, name='tavily_search_results_json'), AIMessage(content=\"Thought: The required information about New York's GDP is provided in the search results. In 2022, New York had a Gross State Product (GSP) of $2.053 trillion.\"), AIMessage(content='The Gross Domestic Product (GDP) of New York in 2022 was $2.053 trillion.')]}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for step in chain.stream([HumanMessage(content=\"What's the GDP of New York?\")]):\n",
    "    print(step)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b96efd08-5314-44f0-a694-3073b638adad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Gross Domestic Product (GDP) of New York in 2022 was $2.053 trillion.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step[END][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c65ef5-b4b2-4ab2-8c78-a551da7819b9",
   "metadata": {},
   "source": [
    "#### Multi-hop question\n",
    "\n",
    "This question requires that the agent perform multiple searches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0b3a0916-d8ca-4092-b91c-d9e2b05259d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': [FunctionMessage(content=\"[{'url': 'https://a-z-animals.com/blog/discover-the-worlds-oldest-parrot/', 'content': 'How Old Is the World’s Oldest Parrot?  Discover the World’s Oldest Parrot Advertisement  of debate, so we’ll detail some other parrots whose lifespans may be longer but are hard to verify their exact age.  Comparing Parrots’ Lifespans to Other BirdsSep 8, 2023 — Sep 8, 2023The oldest parrot on record is Cookie, a pink cockatoo that survived to the age of 83 and survived his entire life at the Brookfield Zoo.'}]\", additional_kwargs={'idx': 0}, name='tavily_search_results_json'), FunctionMessage(content=\"HTTPError('502 Server Error: Bad Gateway for url: https://api.tavily.com/search')\", additional_kwargs={'idx': 1}, name='tavily_search_results_json'), FunctionMessage(content='join', additional_kwargs={'idx': 2}, name='join')]}\n",
      "---\n",
      "{'join': [AIMessage(content='Thought: The oldest parrot on record is Cookie, a pink cockatoo, who lived to be 83 years old. However, there was an error fetching additional search results to compare this age to the average lifespan of parrots.'), SystemMessage(content='Context from last attempt: I found the age of the oldest parrot, Cookie, who lived to be 83 years old. However, I need to search again to find the average lifespan of parrots to complete the comparison.')]}\n",
      "---\n",
      "{'plan_and_schedule': [FunctionMessage(content='[{\\'url\\': \\'https://www.turlockvet.com/site/blog/2023/07/15/parrot-lifespan--how-long-pet-parrots-live\\', \\'content\\': \"Parrot Lifespan  the lifespan of a parrot?\\'.  Parrot Lifespan: How Long Do Pet Parrots Live?  how long they actually live and what you should know about owning a parrot.Jul 15, 2023 — Jul 15, 2023Generally, the average lifespan of smaller species of parrots such as Budgies and Cockatiels is about 5 - 15 years, while larger parrots such as\\\\xa0...\"}]', additional_kwargs={'idx': 3}, name='tavily_search_results_json')]}\n",
      "---\n",
      "{'join': [AIMessage(content=\"Thought: I have found that the oldest parrot on record, Cookie, lived to be 83 years old. Additionally, I've found that the average lifespan of parrots varies by species, with smaller species like Budgies and Cockatiels living between 5-15 years, and larger parrots potentially living longer. This allows me to compare Cookie's age to the average lifespan of smaller parrot species.\"), AIMessage(content=\"The oldest parrot on record is Cookie, a pink cockatoo, who lived to be 83 years old. Compared to the average lifespan of smaller parrot species such as Budgies and Cockatiels, which is about 5-15 years, Cookie lived significantly longer. The average lifespan of larger parrot species wasn't specified, but it's implied that larger parrots may live longer than smaller species, yet likely still much less than 83 years.\")]}\n",
      "---\n",
      "{'__end__': [HumanMessage(content=\"What's the oldest parrot alive, and how much longer is that than the average?\"), FunctionMessage(content=\"[{'url': 'https://a-z-animals.com/blog/discover-the-worlds-oldest-parrot/', 'content': 'How Old Is the World’s Oldest Parrot?  Discover the World’s Oldest Parrot Advertisement  of debate, so we’ll detail some other parrots whose lifespans may be longer but are hard to verify their exact age.  Comparing Parrots’ Lifespans to Other BirdsSep 8, 2023 — Sep 8, 2023The oldest parrot on record is Cookie, a pink cockatoo that survived to the age of 83 and survived his entire life at the Brookfield Zoo.'}]\", additional_kwargs={'idx': 0}, name='tavily_search_results_json'), FunctionMessage(content=\"HTTPError('502 Server Error: Bad Gateway for url: https://api.tavily.com/search')\", additional_kwargs={'idx': 1}, name='tavily_search_results_json'), FunctionMessage(content='join', additional_kwargs={'idx': 2}, name='join'), AIMessage(content='Thought: The oldest parrot on record is Cookie, a pink cockatoo, who lived to be 83 years old. However, there was an error fetching additional search results to compare this age to the average lifespan of parrots.'), SystemMessage(content='Context from last attempt: I found the age of the oldest parrot, Cookie, who lived to be 83 years old. However, I need to search again to find the average lifespan of parrots to complete the comparison. - Begin counting at : 3'), FunctionMessage(content='[{\\'url\\': \\'https://www.turlockvet.com/site/blog/2023/07/15/parrot-lifespan--how-long-pet-parrots-live\\', \\'content\\': \"Parrot Lifespan  the lifespan of a parrot?\\'.  Parrot Lifespan: How Long Do Pet Parrots Live?  how long they actually live and what you should know about owning a parrot.Jul 15, 2023 — Jul 15, 2023Generally, the average lifespan of smaller species of parrots such as Budgies and Cockatiels is about 5 - 15 years, while larger parrots such as\\\\xa0...\"}]', additional_kwargs={'idx': 3}, name='tavily_search_results_json'), AIMessage(content=\"Thought: I have found that the oldest parrot on record, Cookie, lived to be 83 years old. Additionally, I've found that the average lifespan of parrots varies by species, with smaller species like Budgies and Cockatiels living between 5-15 years, and larger parrots potentially living longer. This allows me to compare Cookie's age to the average lifespan of smaller parrot species.\"), AIMessage(content=\"The oldest parrot on record is Cookie, a pink cockatoo, who lived to be 83 years old. Compared to the average lifespan of smaller parrot species such as Budgies and Cockatiels, which is about 5-15 years, Cookie lived significantly longer. The average lifespan of larger parrot species wasn't specified, but it's implied that larger parrots may live longer than smaller species, yet likely still much less than 83 years.\")]}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "steps = chain.stream(\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=\"What's the oldest parrot alive, and how much longer is that than the average?\"\n",
    "        )\n",
    "    ],\n",
    "    {\n",
    "        \"recursion_limit\": 100,\n",
    "    },\n",
    ")\n",
    "for step in steps:\n",
    "    print(step)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6c65c414-7668-4fdf-ba97-f42f659b1317",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The oldest parrot on record is Cookie, a pink cockatoo, who lived to be 83 years old. Compared to the average lifespan of smaller parrot species such as Budgies and Cockatiels, which is about 5-15 years, Cookie lived significantly longer. The average lifespan of larger parrot species wasn't specified, but it's implied that larger parrots may live longer than smaller species, yet likely still much less than 83 years.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step[END][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b859bc7-1a85-4d35-b57b-f67c87282403",
   "metadata": {},
   "source": [
    "#### Multi-step  math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "38d3ea91-59ba-4267-8060-ed75bbc840c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': [FunctionMessage(content='3307.0', additional_kwargs={'idx': 1}, name='math'), FunctionMessage(content='7.565011820330969', additional_kwargs={'idx': 2}, name='math'), FunctionMessage(content='3314.565011820331', additional_kwargs={'idx': 3}, name='math'), FunctionMessage(content='join', additional_kwargs={'idx': 4}, name='join')]}\n",
      "{'join': [AIMessage(content=\"Thought: The calculations for each part of the user's question have been successfully completed. The first calculation resulted in 3307.0, the second in 7.565011820330969, and the sum of those two values was correctly found to be 3314.565011820331.\"), AIMessage(content='The result of ((3*(4+5)/0.5)+3245) + 8 is 3307.0, the result of 32/4.23 is approximately 7.565, and the sum of those two values is approximately 3314.565.')]}\n",
      "{'__end__': [HumanMessage(content=\"What's ((3*(4+5)/0.5)+3245) + 8? What's 32/4.23? What's the sum of those two values?\"), FunctionMessage(content='3307.0', additional_kwargs={'idx': 1}, name='math'), FunctionMessage(content='7.565011820330969', additional_kwargs={'idx': 2}, name='math'), FunctionMessage(content='3314.565011820331', additional_kwargs={'idx': 3}, name='math'), FunctionMessage(content='join', additional_kwargs={'idx': 4}, name='join'), AIMessage(content=\"Thought: The calculations for each part of the user's question have been successfully completed. The first calculation resulted in 3307.0, the second in 7.565011820330969, and the sum of those two values was correctly found to be 3314.565011820331.\"), AIMessage(content='The result of ((3*(4+5)/0.5)+3245) + 8 is 3307.0, the result of 32/4.23 is approximately 7.565, and the sum of those two values is approximately 3314.565.')]}\n"
     ]
    }
   ],
   "source": [
    "for step in chain.stream(\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=\"What's ((3*(4+5)/0.5)+3245) + 8? What's 32/4.23? What's the sum of those two values?\"\n",
    "        )\n",
    "    ]\n",
    "):\n",
    "    print(step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a6cf5fe0-f178-4197-950f-257711bff8d2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The result of ((3*(4+5)/0.5)+3245) + 8 is 3307.0, the result of 32/4.23 is approximately 7.565, and the sum of those two values is approximately 3314.565.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step[END][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c647d5f3-5e00-4449-9cec-5a9f438c9cff",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Congrats on building your first LLMCompiler agent! I'll leave you with some known limitations to the implementation above:\n",
    "\n",
    "1. The planner output parsing format is fragile if your function requires more than 1 or 2 arguments. We could make it more robust by using streaming tool calling.\n",
    "2. Variable substitution is fragile in the example above. It could be made more robust by using a fine-tuned model and a more robust syntax (using e.g., Lark or a tool calling schema)\n",
    "3. The state can grow quite long if you require multiple re-planning runs. To handle, you could add a message compressor once you go above a certain token limit.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "431217e6-4c00-409f-a2bd-40ebff902489",
   "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
}