{
 "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",
    "## Setup\n",
    "\n",
    "First, let's install the required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "16bd5497-35ad-44f2-94d9-19ff39a5ffed",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U --quiet langchain_openai langsmith langgraph langchain numexpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "abbd6948-e9a3-47ca-89c7-7ac2fc5eca8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\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",
    "_get_pass(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d499dad8",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a61b48ee-8c6f-4863-913a-676f659287de",
   "metadata": {},
   "source": [
    "## Define 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/v0.2/docs/integrations/tools/ddg/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "e7476bb2-1a51-42f6-b7ae-82a0300bbf84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_openai import ChatOpenAI\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": [
    "## 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": 78,
   "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 import hub\n",
    "from langchain_core.language_models import BaseChatModel\n",
    "from langchain_core.messages import (\n",
    "    BaseMessage,\n",
    "    FunctionMessage,\n",
    "    HumanMessage,\n",
    "    SystemMessage,\n",
    ")\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnableBranch\n",
    "from langchain_core.tools import BaseTool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from output_parser import LLMCompilerPlanParser, Task\n",
    "\n",
    "prompt = hub.pull(\"wfh/llm-compiler\")\n",
    "print(prompt.pretty_print())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "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": 80,
   "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": 81,
   "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]]) -> 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 0x14e1049a0> {'problem': 'x^3', 'context': ['$1']}\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": [
    "## 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": 82,
   "id": "c1fbafdd-42d4-4575-8466-e5951cee71f4",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "import re\n",
    "import time\n",
    "from concurrent.futures import ThreadPoolExecutor, wait\n",
    "from typing import Any, Dict, Iterable, List, Union\n",
    "\n",
    "from langchain_core.runnables import (\n",
    "    chain as as_runnable,\n",
    ")\n",
    "from typing_extensions import TypedDict\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}, tool_call_id = k\n",
    "        )\n",
    "        for k, (name, task_args, obs) in new_observations.items()\n",
    "    ]\n",
    "    return tool_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "052f6b16-103a-40e9-94dd-8fcc37e77ba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "\n",
    "@as_runnable\n",
    "def plan_and_schedule(state):\n",
    "    messages = state[\"messages\"]\n",
    "    tasks = planner.stream(messages)\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",
    "    )\n",
    "    return {\"messages\": 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": 84,
   "id": "55142257-2674-4a47-988e-0d2810917329",
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_messages = plan_and_schedule.invoke({\"messages\":[HumanMessage(content=example_question)]})['messages']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "a98e0525-2fcf-4fa1-baf6-79858bb8a6bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[FunctionMessage(content=\"[{'url': 'https://www.wunderground.com/weather/us/ca/san-francisco', 'content': 'Current Weather for Popular Cities . San Francisco, CA 82 ° F Sunny; Manhattan, NY warning 84 ° F Sunny; Schiller Park, IL (60176) warning 97 ° F Mostly Cloudy; Boston, MA warning 74 ° F ...'}]\", additional_kwargs={'idx': 1, 'args': {'query': 'current temperature in San Francisco'}}, name='tavily_search_results_json', tool_call_id=1),\n",
       " FunctionMessage(content='551368', additional_kwargs={'idx': 2, 'args': {'problem': 'x ** 3', 'context': ['$1']}}, name='math', tool_call_id=2),\n",
       " FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tool_messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563d5311-55f0-4ca1-afbd-01fd970cf3e3",
   "metadata": {},
   "source": [
    "## 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": 86,
   "id": "942dab42-ad42-4ba2-90d5-49edbe4fae68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\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 = joiner_prompt | llm.with_structured_output(JoinOutputs)"
   ]
  },
  {
   "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": 87,
   "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 {\"messages\": response + [\n",
    "            SystemMessage(\n",
    "                content=f\"Context from last attempt: {decision.action.feedback}\"\n",
    "            )\n",
    "        ]\n",
    "        }\n",
    "    else:\n",
    "        return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n",
    "\n",
    "\n",
    "def select_recent_messages(state) -> dict:\n",
    "    messages = state[\"messages\"]\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": 88,
   "id": "1e49d4b1-8266-4520-a566-1448b1c31c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_messages = [HumanMessage(content=example_question)] + tool_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "31854dfd-b82f-4c24-9b58-6bae66777909",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [AIMessage(content=\"Thought: We have the current temperature in San Francisco (82 °F) and have calculated the temperature raised to the 3rd power (551368). Therefore, we can provide an answer to the user's question.\"),\n",
       "  AIMessage(content='The temperature in San Francisco raised to the 3rd power is 551368.')]}"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joiner.invoke({\"messages\":input_messages})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b099e5ee-2c23-47d9-9387-0f64e02627d3",
   "metadata": {},
   "source": [
    "## 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": 90,
   "id": "768b5f11-e3d2-47be-8143-a7dcd8765243",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.graph.message import add_messages\n",
    "from typing import Annotated\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "graph_builder = StateGraph(State)\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):\n",
    "    messages = state[\"messages\"]\n",
    "    if isinstance(messages[-1], AIMessage):\n",
    "        return END\n",
    "    return \"plan_and_schedule\"\n",
    "\n",
    "\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"join\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    ")\n",
    "graph_builder.add_edge(START, \"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": 91,
   "id": "5bc4584a-e31c-4065-805e-76a6db30676a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': {'messages': [FunctionMessage(content=\"[{'url': 'https://www.investopedia.com/articles/investing/011516/new-yorks-economy-6-industries-driving-gdp-growth.asp', 'content': 'The manufacturing sector is a leader in railroad rolling stock, as many of the earliest railroads were financed or founded in New York; garments, as New York City is the fashion capital of the U.S.; elevator parts; glass; and many other products.\\\\n Educational Services\\\\nThough not typically thought of as a leading industry, the educational sector in New York nonetheless has a substantial impact on the state and its residents, and in attracting new talent that eventually enters the New York business scene. New York has seen a large uptick in college attendees, both young and old, over the 21st century, and an increasing number of new employees in other New York sectors were educated in the state. New York City is the leading job hub for banking, finance, and communication in the U.S. New York is also a major manufacturing center and shipping port, and it has a thriving technological sector.\\\\n The state of New York has the third-largest economy in the United States with a gross domestic product (GDP) of $1.7 trillion, trailing only Texas and California.'}]\", additional_kwargs={'idx': 1, 'args': {'query': 'GDP of New York'}}, name='tavily_search_results_json', tool_call_id=1)]}}\n",
      "---\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: The information required to answer the user's question has been found. The GDP of New York is mentioned as $1.7 trillion, making it the third-largest economy in the United States.\", id='d656a605-e4c4-470d-9b29-31794f298a71'), AIMessage(content='The GDP of New York is $1.7 trillion, making it the third-largest economy in the United States.', id='5135758e-d01e-4360-bb6a-31025b723d8c')]}}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for step in chain.stream(\n",
    "    {\"messages\": [HumanMessage(content=\"What's the GDP of New York?\")]}\n",
    "):\n",
    "    print(step)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "b96efd08-5314-44f0-a694-3073b638adad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The GDP of New York is $1.7 trillion, making it the third-largest economy in the United States.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step['join']['messages'][-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": 93,
   "id": "0b3a0916-d8ca-4092-b91c-d9e2b05259d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': {'messages': [FunctionMessage(content='[{\\'url\\': \\'https://en.wikipedia.org/wiki/Cookie_(cockatoo)\\', \\'content\\': \\'He was one of the longest-lived birds on record[4] and was recognised by the Guinness World Records as the oldest living parrot in the world.[5]\\\\nThe next-oldest pink cockatoo to be found in a zoological setting was a 31-year-old female bird located at Paradise Wildlife Sanctuary, England.[3] Information published by the World Parrot Trust states longevity for Cookie\\\\\\'s species in captivity is on average 40–60 years.[6]\\\\nLife[edit]\\\\nCookie was Brookfield Zoo\\\\\\'s oldest resident and the last surviving member of the animal collection from the time of the zoo\\\\\\'s opening in 1934, having arrived from Taronga Zoo of Sydney, New South Wales, Australia, in the same year and judged to be one year old at the time.[7]\\\\nIn the 1950s an attempt was made to introduce Cookie to a female pink cockatoo, but Cookie rejected her as \"she was not nice to him\".[8]\\\\n In 2007, Cookie was diagnosed with, and placed on medication and nutritional supplements for, osteoarthritis and osteoporosis\\\\xa0– medical conditions which occur commonly in aging animals and humans alike,[7] although it is believed that the latter may also have been brought on as a result of being fed a seed-only diet for the first 40 years of his life, in the years before the dietary requirements of his species were fully understood.[9]\\\\nCookie was \"retired\" from exhibition at the zoo in 2009 (following a few months of weekend-only appearances) in order to preserve his health, after it was noticed by staff that his appetite, demeanor and stress levels improved markedly when not on public display. age.[11] A memorial at the zoo was unveiled in September 2017.[12]\\\\nIn 2020, Cookie became the subject of a poetry collection by Barbara Gregorich entitled Cookie the Cockatoo: Everything Changes.[13]\\\\nSee also[edit]\\\\nReferences[edit]\\\\nExternal links[edit] He was believed to be the oldest member of his species alive in captivity, at the age of 82 in June 2015,[1][2] having significantly exceeded the average lifespan for his kind.[3] He was moved to a permanent residence in the keepers\\\\\\' office of the zoo\\\\\\'s Perching Bird House, although he made occasional appearances for special events, such as his birthday celebration, which was held each June.[3]\\'}]', additional_kwargs={'idx': 1, 'args': {'query': 'oldest parrot alive'}}, name='tavily_search_results_json', tool_call_id=1), FunctionMessage(content='[{\\'url\\': \\'https://www.thesprucepets.com/how-long-do-parrots-and-other-pet-birds-live-1238433\\', \\'content\\': \"It\\'s possible that a pet bird can outlive its owners\\\\nThe Spruce / Adrienne Legault\\\\nParrots and other birds can live up to 10 to 50 years or more depending on the type and the conditions they live in. They vary in size from small birds that can fit in the palm of your hand to large birds the size of a cat and their lifespans are just as variable.\\\\n Also, for birds who live longer some owners have to make a plan of where the bird is going in the circumstance the bird outlives the owner.\\\\n In reality, there is a wide range in the age that pet birds might reach and certainly, some will live longer (or shorter amounts of time) than the ages listed.\\\\n Potential owners need to be aware of the longevity of their bird so they can be prepared to provide proper care for them for as long as they live.\\\\n\"}]', additional_kwargs={'idx': 2, 'args': {'query': 'average lifespan of a parrot'}}, name='tavily_search_results_json', tool_call_id=2), FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]}}\n",
      "---\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: We have information on Cookie, the cockatoo, who was recognized as the oldest living parrot at 82 years old in June 2015. This significantly exceeds the average lifespan for his kind, which is stated to be 40-60 years. The second source provides a general lifespan range for parrots and other birds, which is 10-50 years. However, this range varies significantly depending on the species and conditions. Since Cookie's specific lifespan far exceeds the average for his species and falls outside the general range for parrots, we can answer the user's question.\", id='51a280ac-2327-40c5-a27a-c821697d5a4b'), AIMessage(content='The oldest parrot recorded was Cookie, a cockatoo, who lived to be 82 years old in June 2015. This is significantly longer than the average lifespan for his species, which is 40-60 years, and also exceeds the general lifespan range for parrots, which can vary from 10 to 50 years. Therefore, Cookie lived 22 to 42 years longer than the average lifespan for his species.', id='139ecedf-b090-4197-88c0-0fa39883b392')]}}\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "steps = chain.stream({\"messages\":\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=\"What's the oldest parrot alive, and how much longer is that than the average?\"\n",
    "        )\n",
    "    ]\n",
    "    },\n",
    "    {\n",
    "        \"recursion_limit\": 100,\n",
    "    },\n",
    ")\n",
    "for step in steps:\n",
    "    print(step)\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "6c65c414-7668-4fdf-ba97-f42f659b1317",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The oldest parrot recorded was Cookie, a cockatoo, who lived to be 82 years old in June 2015. This is significantly longer than the average lifespan for his species, which is 40-60 years, and also exceeds the general lifespan range for parrots, which can vary from 10 to 50 years. Therefore, Cookie lived 22 to 42 years longer than the average lifespan for his species.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step['join']['messages'][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b859bc7-1a85-4d35-b57b-f67c87282403",
   "metadata": {},
   "source": [
    "### Multi-step  math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "38d3ea91-59ba-4267-8060-ed75bbc840c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': {'messages': [FunctionMessage(content='3307.0', additional_kwargs={'idx': 1, 'args': {'problem': '((3*(4+5)/0.5)+3245) + 8'}}, name='math', tool_call_id=1), FunctionMessage(content='7.565011820330969', additional_kwargs={'idx': 2, 'args': {'problem': '32/4.23'}}, name='math', tool_call_id=2), FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]}}\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: The calculations for both individual questions have been provided: 3307.0 for the first equation and 7.565011820330969 for the second. To answer the user's final question, we need to sum these two values.\", id='96eb85f5-831f-434e-83d8-59deeebce05d'), AIMessage(content='The result of the first calculation is 3307.0, and the result of the second calculation is approximately 7.57. The sum of those two values is approximately 3314.57.', id='671a1a08-4725-4f98-997a-848815d61aa5')]}}\n"
     ]
    }
   ],
   "source": [
    "for step in chain.stream({\"messages\":\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": 97,
   "id": "a6cf5fe0-f178-4197-950f-257711bff8d2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The result of the first calculation is 3307.0, and the result of the second calculation is approximately 7.57. The sum of those two values is approximately 3314.57.\n"
     ]
    }
   ],
   "source": [
    "# Final answer\n",
    "print(step['join']['messages'][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9487866",
   "metadata": {},
   "source": [
    "### Complex Replanning Example\n",
    "\n",
    "This question is likely to prompt the Replan functionality, but it may need to be run multiple times to see this in action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391d6931",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'plan_and_schedule': {'messages': [FunctionMessage(content=\"[{'url': 'https://www.timeanddate.com/weather/japan/tokyo', 'content': '88 / 84 °F. 13. 87 / 82 °F. 14. 84 / 80 °F. Detailed forecast for 14 days. Need some help? Current weather in Tokyo and forecast for today, tomorrow, and next 14 days.'}]\", additional_kwargs={'idx': 1, 'args': {'query': 'current temperature in Tokyo'}}, name='tavily_search_results_json', tool_call_id=1), FunctionMessage(content='join', additional_kwargs={'idx': 2, 'args': ()}, name='join', tool_call_id=2)]}}\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides the current temperature in Tokyo but does not explicitly state which temperature (88 / 84 °F) corresponds to the current condition. It seems to be a range, possibly the day's high and low. Without a clear indication of the exact current temperature, it's challenging to provide a precise flashcard summary.\", id='8ef2a131-69db-4180-a76e-fd9d6f4037c1'), SystemMessage(content='Context from last attempt: The information provided does not explicitly state the current temperature in Tokyo; it provides a temperature range without specifying which is the current temperature. Need to find a source that gives the exact current temperature in Tokyo for a precise flashcard summary.', id='f5bd752c-b068-459a-8d9e-bd1f1b5fa4fe')]}}\n",
      "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]}}\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='3cc41891-4f47-4453-8edf-b989926ab25e'), SystemMessage(content='Context from last attempt: The search did not provide an exact current temperature for Tokyo, making it impossible to create a precise flashcard. A source that explicitly states the current temperature is needed for an accurate response.', id='96290b41-a4c4-4ab5-829a-89cc31dfe6c8')]}}\n",
      "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 4, 'args': ()}, name='join', tool_call_id=4)]}}\n",
      "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='4724b242-ddb8-47e6-b235-de25de54fe45'), AIMessage(content='I was unable to find the exact current temperature in Tokyo. However, the temperature range for today in Tokyo is between 88°F and 84°F. For the most accurate and up-to-date temperature, I recommend checking a reliable weather forecasting website or app.', id='40e29a47-a001-4f65-a18f-65c2931d1ae5')]}}\n"
     ]
    }
   ],
   "source": [
    "for step in chain.stream({\"messages\":\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=\"Find the current temperature in Tokyo, then, respond with a flashcard summarizing this information\"\n",
    "        )\n",
    "    ]}\n",
    "):\n",
    "    print(step)"
   ]
  },
  {
   "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"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}