langgraph/examples/advanced_agents/LLMCompiler.ipynb

Implementing LLMCompiler using LangGraph

By Kim, et. al 🔗

LLMCompiler is an agent architecture intented on speeding up the latency of agentic tasks via fast, parallel tool execution. It has 3 main components:

1. Planner: generate a DAG of tasks.
2. Task Fetching Unit: schedules and executes the tasks
3. Joiner: Responds to the user or triggers a second plan

This notebook walks through each component and shows how to wire them together using LangGraph.

Part 1: Planner

Largely adapted from the original source code.

Output Parser

Parses task lists in the following form:

1. tool_1("arg1", 3.5, ...)
Thought: I then want to find out Y by using tool_2
2. tool_2("", ${1})'
3. join()<END_OF_PLAN>"

The "Thought" lines are optional. The ${#} placeholders are variables. These are used to route tool (task) outputs to other tools.

In [1]:

import json
import re
from typing import Any, Dict, List, Optional, Sequence, Union

from langchain.agents.agent import AgentOutputParser
from langchain.schema import OutputParserException
from langchain_core.tools import BaseTool

THOUGHT_PATTERN = r"Thought: ([^\n]*)"
# $1 or ${1} -> 1
ID_PATTERN = r"\$\{?(\d+)\}?"
END_OF_PLAN = "<END_OF_PLAN>"


class ActionParserFSM:
    def __init__(self):
        self.reset()

    def reset(self):
        self.state = "START"
        self.task_index = ""
        self.action = ""
        self.comment = ""
        self.bracket_count = 0
        self.actions = []

    def parse(self, text: str):
        for char in text:
            action = self.process_char(char)
            if action:
                yield action
        action = self.save_action()
        if action:
            yield action

    def process_char(self, char: str) -> Optional[dict]:
        action = None
        if self.state == "START":
            if char.isdigit():
                self.state = "NUMBER"
                self.task_index += char
            elif char == "\n":
                self.reset()
        elif self.state == "NUMBER":
            if char == ".":
                self.state = "ACTION"
            elif char.isdigit():
                self.task_index += char
            else:
                self.reset()
        elif self.state == "ACTION":
            if char == "{":
                self.bracket_count += 1
            elif char == "}":
                self.bracket_count -= 1
                if self.bracket_count == 0:
                    self.state = "COMMENT"
            self.action += char
        elif self.state == "COMMENT":
            if char == "\n":
                action = self.save_action()
                self.reset()
            else:
                self.comment += char
        return action

    def save_action(self):
        if self.task_index and self.action:
            parsed_action = json.loads(self.action.strip())
            tool_name, args = next(iter(parsed_action.items()))
            return {
                "task_index": int(self.task_index),
                "tool_name": tool_name,
                "args": args,
            }


class LLMCompilerPlanParser(AgentOutputParser, extra="allow"):
    """Planning output parser."""

    def __init__(self, tools: Sequence[BaseTool], **kwargs):
        super().__init__(**kwargs)
        self.tools = tools

    def parse(self, text: str) -> list[str]:
        parser = ActionParserFSM()
        graph_dict = {}
        for task in parser.parse(text):
            idx = int(task["task_index"])

            task = instantiate_task(
                tools=self.tools,
                idx=idx,
                tool_name=task["tool_name"],
                args=task["args"],
            )

            graph_dict[idx] = task
            if task["tool"] == "join":
                break

        return graph_dict


### Helper functions


def default_dependency_rule(idx, args: str):
    matches = re.findall(ID_PATTERN, args)
    numbers = [int(match) for match in matches]
    return idx in numbers


def _get_dependencies_from_graph(
    idx: int, tool_name: str, args: Sequence[Any]
) -> dict[str, list[str]]:
    """Get dependencies from a graph."""
    if tool_name == "join":
        return list(range(1, idx))
    return [i for i in range(1, idx) if default_dependency_rule(i, str(args))]


def instantiate_task(
    tools: Sequence[BaseTool],
    idx: int,
    tool_name: str,
    args: Union[dict, str, bool, None],
) -> dict:
    dependencies = _get_dependencies_from_graph(idx, tool_name, args)
    if tool_name == "join":
        tool = "join"
    else:
        try:
            tool = tools[[tool.name for tool in tools].index(tool_name)]
        except ValueError as e:
            raise OutputParserException(f"Tool {tool_name} not found.")
    return dict(
        tool=tool,
        args=args,
        dependencies=dependencies,
    )

Planner Code

This takes the input and outputs a plan.

In [141]:

from langchain.chat_models.base import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableBranch
from langchain_core.tools import BaseTool

END_OF_PLAN = "<END_OF_PLAN>"

JOINER_FINISH = "Finish"
JOINER_REPLAN = "Replan"


JOIN_DESCRIPTION = (
    "join():\n"
    " - Collects and combines results from prior actions.\n"
    " - A 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. "
)

planner_prompt_tmpl_str = (
    "Given a user query, create a plan to solve it with the utmost parallelizability. "
    "Each plan should comprise an action from the following {num_tools} types:\n"
    "{tool_descriptions}"
    f"\n{{num_toolsp1}}. {JOIN_DESCRIPTION}"
    "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.\n"
    " - Provide actions ONLY in json form, with the single key being the action name and the value being its arguments. Do not write python code. \n"
    " - Each action line must start with 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"
    f" - 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"
    "{replan}"
    "{examples}"
)


def _generate_planner_prompt(
    tools: Sequence[BaseTool],
    example_prompt=str,
):
    tool_descriptions = "\n".join(
        f"{i+1}. {tool.name}: {tool.description}\n\tInput schema: {tool.args}"
        for i, tool in enumerate(tools)
    )
    planner_prompt_template = ChatPromptTemplate.from_messages(
        [("system", planner_prompt_tmpl_str), ("user", "Question: {input}{context}")]
    ).partial(
        tool_descriptions=tool_descriptions,
        examples="Here are some examples:\n\n" + example_prompt
        if example_prompt
        else "",
        num_tools=len(tools),
        num_toolsp1=len(tools) + 1,
    )

    return planner_prompt_template


def create_planner(
    llm: BaseChatModel,
    example_prompt: str,
    tools: Sequence[BaseTool],
    stop: Optional[list[str]] = None,
):
    og_planner_prompt = _generate_planner_prompt(tools, example_prompt).partial(
        replan="",
        context="",
    )
    replanner_prompt = _generate_planner_prompt(tools, example_prompt).partial(
        replan=' - You are given "Previous Plan" which is the plan that the previous agent created along with the execution results '
        "(given as Observation) of each plan and a general thought (given as Thought) about the executed results."
        'You MUST use these information to create the next plan under "Current Plan".\n'
        ' - When starting the Current Plan, you should start with "Thought" that outlines the strategy for the next plan.\n'
        " - In the Current Plan, you should NEVER repeat the actions that are already executed in the Previous Plan.\n"
        " - You must continue the task index from the end of the previous one. Do not repeat task indices."
    )
    bound_llm = llm.bind(stop=stop)
    return (
        RunnableBranch(
            ((lambda x: x.get("context") is not None), replanner_prompt),
            og_planner_prompt,
        )
        | bound_llm
        | LLMCompilerPlanParser(tools=tools)
    )

Example usage

Here's an example usage of the planner module

In [131]:

from langchain.tools import tool
from langchain_openai import ChatOpenAI


@tool
def get_user_id(first_name: str, last_name: Optional[str] = None):
    """Query the user IDs of everyone with the provided name."""
    return 4


@tool
def get_scores(class_name: str, user_id: int):
    """Query the class registry for grades of the provided user ID."""
    return "A+"


examples = (
    "Question: What's the user ID for Johnny Drop Tables?\n"
    '1. {"get_user_id": {"first_name": "Johnny", "last_name":"Drop Tables"}}\n'
    f'2. {{"join": null}}{END_OF_PLAN}\n'
    "###\n"
    "\n"
    "Question: What was Eric Zhang's score in Calc?\n"
    '1. {"get_user_id": {"first_name": "Eric", "last_name":"Zhang"}}\n'
    '2. {"get_scores": {"class_name": "calc", "user_id": "$1"}}\n'
    f'3. {{"join": null}}{END_OF_PLAN}\n'
    "###\n"
    "\n"
)

planner = create_planner(
    ChatOpenAI(model="gpt-3.5-turbo"),
    example_prompt=examples,
    tools=[get_user_id, get_scores],
)

In [132]:

tasks = planner.invoke(
    {"input": "What are the Calc BC grades for Sam and Will Van Damm?"}
)
tasks

Out [132]:

{1: {'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: Optional[str] = None) - Query the user IDs of everyone with the provided name.', args_schema=<class 'pydantic.main.get_user_idSchemaSchema'>, func=<function get_user_id at 0x130eb9a80>),
  'args': {'first_name': 'Sam', 'last_name': 'Van Damm'},
  'dependencies': []},
 2: {'tool': StructuredTool(name='get_user_id', description='get_user_id(first_name: str, last_name: Optional[str] = None) - Query the user IDs of everyone with the provided name.', args_schema=<class 'pydantic.main.get_user_idSchemaSchema'>, func=<function get_user_id at 0x130eb9a80>),
  'args': {'first_name': 'Will', 'last_name': 'Van Damm'},
  'dependencies': []},
 3: {'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) - Query the class registry for grades of the provided user ID.', args_schema=<class 'pydantic.main.get_scoresSchemaSchema'>, func=<function get_scores at 0x130eb9d00>),
  'args': {'class_name': 'Calc BC', 'user_id': '$1'},
  'dependencies': [1]},
 4: {'tool': StructuredTool(name='get_scores', description='get_scores(class_name: str, user_id: int) - Query the class registry for grades of the provided user ID.', args_schema=<class 'pydantic.main.get_scoresSchemaSchema'>, func=<function get_scores at 0x130eb9d00>),
  'args': {'class_name': 'Calc BC', 'user_id': '$2'},
  'dependencies': [2]},
 5: {'tool': 'join', 'args': None, 'dependencies': [1, 2, 3, 4]}}

2. Task Fetching Unit

This component schedules the tasks. In the paper, it's kept separate from the "executor", but here we create a single DAG to be executed by LangGraph.

Basic idea is that, given a list of dicts of the form:

{
    tool: BaseTool,
    dependencies: number[],
}

1. Create a topological sort of the tasks
2. Execute them on the previous step's output, ensuring to perform variable substitution where appropriate

In [184]:

import functools

from langchain_core.runnables import (
    RunnableLambda,
    RunnableParallel,
    RunnablePassthrough,
)


def _sort_tasks(data):
    if not data:
        return []
    sorted_tasks = []
    # Remove tasks already completed
    min_idx = min([int(k) for k in data])
    data = {
        int(k): {
            **v,
            "dependencies": [dep for dep in v["dependencies"] if dep >= min_idx],
        }
        for k, v in data.items()
    }
    while data:
        no_deps = {k: v for k, v in data.items() if not v["dependencies"]}
        if not no_deps:
            raise ValueError("We seem to have run into a circular dependency.")

        sorted_tasks.append(no_deps)
        data = {
            k: {
                **v,
                "dependencies": [d for d in v["dependencies"] if d not in no_deps],
            }
            for k, v in data.items()
            if k not in no_deps
        }
    return sorted_tasks


def _resolve_arg(x: dict, arg: Union[str, Any]):
    if isinstance(arg, str) and arg.startswith("$"):
        try:
            return x[f"task_{arg[1:]}"]
        except:
            if arg.endswith(".output"):
                return x[f"task_{arg[1:-7]}"]
            raise

    else:
        return arg


def _execute_task(x, task):
    tool_to_use = task["tool"]
    args = task["args"]
    if isinstance(args, str):
        resolved_args = _resolve_arg(x, args)
    elif isinstance(args, dict):
        resolved_args = {key: _resolve_arg(x, val) for key, val in args.items()}
    else:
        # This will likely fail
        resolved_args = args
    try:
        return tool_to_use.invoke(resolved_args)
    except Exception as e:
        return (
            f"ERROR(Failed to call tool {tool_to_use} with args {tool_to_use}."
            + f" Args resolved to {resolved_args}. Error: {repr(e)})"
        )


def construct_dag(tasks):
    sorted_tasks = _sort_tasks(tasks)
    chain = None
    for idx, task_group in enumerate(sorted_tasks):
        if len(task_group) == 1 and next(iter(task_group.values()))["tool"] == "join":
            # TODO: actually join the values
            step = lambda x: {"join": x}
        else:
            # Cascade all results forward
            constructor = (
                RunnableParallel if chain is None else RunnablePassthrough.assign
            )
            task_dict = {}
            for idx, task in task_group.items():
                task_dict[f"task_{idx}"] = RunnableLambda(
                    functools.partial(_execute_task, task=task)
                ).with_config(run_name=f"task_{idx}")

            step = constructor(**task_dict).with_config(run_name=f"TaskGroup{idx}")
        if chain is None:
            chain = step
        else:
            chain |= step

    if chain is not None:
        return chain | RunnablePassthrough.assign(tasks=lambda _: tasks)
    return chain

In [185]:

graph = construct_dag(tasks)
graph.get_graph().print_ascii()

                      +------------------------------+                       
                      | Parallel<task_1,task_2>Input |                       
                      +------------------------------+                       
                             ***            ***                              
                           **                  **                            
                         **                      **                          
               +-------------+               +-------------+                 
               | Lambda(...) |               | Lambda(...) |                 
               +-------------+               +-------------+                 
                             ***            ***                              
                                **        **                                 
                                  **    **                                   
                     +-------------------------------+                       
                     | Parallel<task_1,task_2>Output |                       
                     +-------------------------------+                       
                                      *                                      
                                      *                                      
                                      *                                      
                      +------------------------------+                       
                      | Parallel<task_3,task_4>Input |                       
                      +------------------------------+                       
                    *****             *           *****                      
               *****                  *                *****                 
            ***                       *                     ***              
+-------------+               +-------------+               +-------------+  
| Lambda(...) |               | Lambda(...) |               | Passthrough |  
+-------------+*****          +-------------+          *****+-------------+  
                    *****             *           *****                      
                         *****        *      *****                           
                              ***     *   ***                                
                     +-------------------------------+                       
                     | Parallel<task_3,task_4>Output |                       
                     +-------------------------------+                       
                                      *                                      
                                      *                                      
                                      *                                      
                     +-------------------------------+                       
                     | Lambda(lambda x: {'join': x}) |                       
                     +-------------------------------+                       
                                      *                                      
                                      *                                      
                                      *                                      
                          +----------------------+                           
                          | Parallel<tasks>Input |                           
                          +----------------------+                           
                           ***               ***                             
                        ***                     ***                          
                      **                           **                        
      +-------------------------+               +-------------+              
      | Lambda(lambda _: tasks) |               | Passthrough |              
      +-------------------------+               +-------------+              
                           ***               ***                             
                              ***         ***                                
                                 **     **                                   
                         +-----------------------+                           
                         | Parallel<tasks>Output |                           
                         +-----------------------+                           

Example Plan

We still haven't introduced any cycles in our computation graph, so this is all easily expressed in LCEL.

In [135]:

chain = planner | construct_dag

In [136]:

example_question = "Did Aliya get a better score than Roger in Geology?"
task_results = chain.invoke({"input": example_question})
task_results["join"]

Out [136]:

{'task_1': 4, 'task_2': 4, 'task_3': 'A+', 'task_4': 'A+'}

Agent Logic

So now we have the planning and initial execution done. We need a component to process these outputs and either:

1. Respond with the correct answer.
2. Loop with a new plan.

The paper calls this the "joiner".

In [216]:

from langchain_core.output_parsers import StrOutputParser
from typing_extensions import TypedDict


def format_task(task, idx):
    tool = task["tool"]
    tool_name = tool if isinstance(tool, str) else tool.name  # Handle join()
    return f"{idx}. {{{tool_name}: {task['args']}}}"


def format_tasks(executor_output: dict):
    tasks = executor_output["tasks"]
    prior_observations = executor_output.get("observations")
    formatted_plan = "\n".join(format_task(task, idx) for idx, task in tasks.items())
    observations = "\n".join(f"{k}: {v}" for k, v in executor_output["join"].items())
    result = f"Original Plan:\n{formatted_plan}\nExecuted plan results:\n{observations}"
    if prior_observations:
        result += f"\nPrevious Results:\n{prior_observations}"
    return result


def _parse_joiner_output(raw_answer: str) -> str:
    thought, answer, is_replan = "", "", False  # default values
    raw_answers = raw_answer.split("\n")
    for ans in raw_answers:
        if ans.startswith("Action:"):
            answer = ans[ans.find("(") + 1 : ans.find(")")]
            is_replan = JOINER_REPLAN in ans
        elif ans.startswith("Thought:"):
            thought = ans.split("Thought:")[1].strip()
    if is_replan:
        return {"thought": thought, "context": answer}
    else:
        return {"thought": thought, "answer": answer}

In [217]:

def create_joiner(prompt, llm):
    return (
        (
            lambda x: {
                **x["plan"],
                "input": x["input"],
                "context": x.get("context"),
                "observations": x.get("observations"),
            }
        )
        | RunnablePassthrough.assign(scratchpad=format_tasks)
        | ChatPromptTemplate.from_messages([("system", prompt), ("user", "{input}")])
        | llm
        | StrOutputParser()
        | _parse_joiner_output
    )

In [218]:

system_prompt = (
    "Solve a question answering task. Here are some guidelines:\n"
    " - In the Assistant Scratchpad, you will be given results of a plan you have executed to answer the user's question.\n"
    " - Thought needs to reason about the question based on the Observations in 1-2 sentences.\n"
    " - Ignore irrelevant action results.\n"
    " - If the required information is present, give a concise but complete and helpful answer to the user's question.\n"
    " - If you are unable to give a satisfactory finishing answer, replan to get the required information."
    " Respond in the following format:\n\n"
    "Thought: <reason about the task results and whether you have sufficient information to answer the question>\n"
    "Action: <action to take>\n"
    "Available actions:\n"
    f" (1) {JOINER_FINISH}(the final answer to return to the user): returns the answer and finishes the task.\n"
    f" (2) {JOINER_REPLAN}(the reasoning and other information that will help you plan again. Can be a line of any length): instructs why we must replan\n\n"
    " Examples:\n"
    "Question: How many users are currently using the new product?\n"
    "...task returns the number 32,000\n"
    "Thought: I find no issue with the original plan, and the results satisfy everything in the user question.\n"
    f"Action: {JOINER_FINISH}(32,000 users currently use the new product)\n###\n"
    "Question: How much cooler is it in NY than SF?\n"
    "...task results show SF is 57 degrees fahrenheit today, and they show in NY it has a high of 32 degrees fahrenheit \n"
    "Thought: I can answer by synthesizing the results.\n"
    f"Action: {JOINER_FINISH}(NY is 25 degrees cooler than SF today, as it has a high of 32 degrees Fahrenheit today, whereas in SF, it is 57 degrees Fahrenheit.)\n###\n"
    "Question: Are the gophers beating the rabbits??\n"
    "...task returns the a score of 7 for rabbits but no other value...\n"
    "Thought: I need the gophers' score to make a final decision.\n"
    f"Action: {JOINER_REPLAN}(The rabbits have a score of 7, but I need the gophers' score.)"
    "\nAssistant Scratchpad:\n{scratchpad}"
)

In [219]:

joiner = create_joiner(system_prompt, ChatOpenAI(model="gpt-3.5-turbo"))
joiner.invoke({"plan": task_results, "input": example_question})

Out [219]:

{'thought': 'The plan has been executed successfully and returned the scores for both Aliya and Roger in Geology. We can compare their scores to determine if Aliya got a better score than Roger.',
 'answer': 'Aliya got an A+ in Geology, while Roger also got an A+. Therefore, they both got the same score in Geology.'}

Construct Agent

Now we have all the required pieces! Let's construct our agent with its tools.

In [220]:

import getpass
import os

os.environ["TAVILY_API_KEY"] = (
    os.environ.get("TAVILY_API_KEY")
    if "TAVILY_API_KEY" in os.environ
    else getpass.getpass("Tavily API Key:")
)
# Then fetch a credentials.json file
# https://developers.google.com/gmail/api/quickstart/python#authorize_credentials_for_a_desktop_application

In [221]:

from operator import add, mul, sub, truediv
from typing import Literal

from langchain_community.agent_toolkits import GmailToolkit
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool


@tool
def calculate(
    arg1: float,
    arg2: float,
    op: Union[Literal["+"], Literal["-"], Literal["*"], Literal["/"]],
):
    """Calculate a mathematical operation on two arguments."""
    resolved_op = {"+": add, "-": sub, "*": mul, "/": truediv}
    return resolved_op[op](arg1, arg2)


tools = [TavilySearchResults(max_results=1), calculate]

In [224]:

from langgraph.graph import END, StateGraph


class GraphState(TypedDict):
    input: str
    plan: Dict
    agent_output: Dict
    observations: Dict
    num_iterations: int
    context: str
    stop_reason: str


MAX_ITERATIONS = 5
workflow = StateGraph(GraphState)

# 1.  Define vertices

planner = create_planner(
    llm=ChatOpenAI(model="gpt-4-1106-preview"),
    # Add more examples to improve reliability
    example_prompt=(
        "Question: What's the capital of Myanmar?\n"
        '1. {"tavily_search_results_json": {"query": "Capital of Myanmar"}}\n'
        f'2. {{"join": null}}{END_OF_PLAN}\n'
        "###\n"
        "\n"
    ),
    tools=tools,
)

plan_and_execute = planner | construct_dag
joiner = create_joiner(system_prompt, ChatOpenAI(model="gpt-4-1106-preview"))


def _reformat_task(idx, task: Union[BaseTool, str]):
    tool = task["tool"]
    tool_name = tool if isinstance(tool, str) else tool.name
    called = {tool_name: task["args"]}
    return f"{idx}. {json.dumps(called)}"


def provide_context(state):
    # Insert a context string for the re-planner.
    # This could alternatively call an LLM to provide additional logic
    context = state["agent_output"]["context"]
    num_iterations = int(state.get("num_iterations") or 1) + 1
    previous_plan = "\n".join(
        [
            _reformat_task(idx, task)
            for idx, task in sorted(state["plan"]["tasks"].items())
        ]
    )
    context_str = (
        f"\n\nPrevious Plan:\n{previous_plan}\n"
        f"{context}\nYou have made {num_iterations}/{MAX_ITERATIONS} attempts thus far."
    )
    observations = state["observations"] or {}
    for task, observation in state["plan"]["join"].items():
        observations[task] = observation
    return {
        "context": context_str,
        "num_iterations": num_iterations,
        "observations": observations,
    }


def add_stop_reason(state):
    num_iterations = int(state.get("num_iterations") or 0)
    if num_iterations >= MAX_ITERATIONS:
        return {"stop_reason": "end_max_iter"}
    if state["agent_output"].get("answer"):
        return {"stop_reason": "answer"}
    return {"stop_reason": None}


# Assign each node to a state variable to update
workflow.add_node("plan_and_execute", RunnablePassthrough.assign(plan=plan_and_execute))
workflow.add_node("join", RunnablePassthrough.assign(agent_output=joiner))
workflow.add_node("provide_context", provide_context)
workflow.add_node("provide_stop_reason", add_stop_reason)


## Define edges

workflow.add_edge("plan_and_execute", "join")
workflow.add_edge("provide_context", "plan_and_execute")
workflow.add_edge("join", "provide_stop_reason")

### This condition determines looping logic


def should_continue(state):
    if state["stop_reason"] is None:
        return "continue"
    return "end"


workflow.add_conditional_edges(
    start_key="provide_stop_reason",
    # Next, we pass in the function that will determine which node is called next.
    condition=should_continue,
    conditional_edge_mapping={
        # If it generates context, we must replan
        "continue": "provide_context",
        # Otherwise we finish.
        "end": END,
    },
)
workflow.set_entry_point("plan_and_execute")
chain = workflow.compile()

Simple question

Let's ask a simple question of the agent.

In [225]:

result = chain.invoke({"input": "What's the GDP of New York?"})
print(result["agent_output"]["answer"])

The GDP of New York in 2022 was about 1.56 trillion U.S. dollars.

Multi-hop question

In [227]:

result = chain.invoke(
    {"input": "How much larger is the GDP of the UK than that of New York?"}
)

---------------------------------------------------------------------------
JSONDecodeError                           Traceback (most recent call last)
Cell In[227], line 1
----> 1 result = chain.invoke(
      2     {"input": "How much larger is the GDP of the UK than that of New York?"}
      3 )

File ~/code/lc/langgraph/langgraph/pregel/__init__.py:492, in Pregel.invoke(self, input, config, output_keys, input_keys, **kwargs)
    482 def invoke(
    483     self,
    484     input: Union[dict[str, Any], Any],
   (...)
    489     **kwargs: Any,
    490 ) -> Union[dict[str, Any], Any]:
    491     latest: Union[dict[str, Any], Any] = None
--> 492     for chunk in self.stream(
    493         input,
    494         config,
    495         output_keys=output_keys if output_keys is not None else self.output,
    496         input_keys=input_keys,
    497         **kwargs,
    498     ):
    499         latest = chunk
    500     return latest

File ~/code/lc/langgraph/langgraph/pregel/__init__.py:528, in Pregel.transform(self, input, config, output_keys, input_keys, **kwargs)
    519 def transform(
    520     self,
    521     input: Iterator[Union[dict[str, Any], Any]],
   (...)
    526     **kwargs: Any,
    527 ) -> Iterator[Union[dict[str, Any], Any]]:
--> 528     for chunk in self._transform_stream_with_config(
    529         input,
    530         self._transform,
    531         config,
    532         output_keys=output_keys,
    533         input_keys=input_keys,
    534         **kwargs,
    535     ):
    536         yield chunk

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1226, in Runnable._transform_stream_with_config(self, input, transformer, config, run_type, **kwargs)
   1224 try:
   1225     while True:
-> 1226         chunk: Output = context.run(next, iterator)  # type: ignore
   1227         yield chunk
   1228         if final_output_supported:

File ~/code/lc/langgraph/langgraph/pregel/__init__.py:313, in Pregel._transform(self, input, run_manager, config, input_keys, output_keys)
    303 done, inflight = concurrent.futures.wait(
    304     [
    305         executor.submit(proc.invoke, input, config)
   (...)
    309     timeout=self.step_timeout,
    310 )
    312 # interrupt on failure or timeout
--> 313 _interrupt_or_proceed(done, inflight, step)
    315 # apply writes to channels
    316 _apply_writes(checkpoint, channels, pending_writes, config, step + 1)

File ~/code/lc/langgraph/langgraph/pregel/__init__.py:611, in _interrupt_or_proceed(done, inflight, step)
    609             inflight.pop().cancel()
    610         # raise the exception
--> 611         raise exc
    612         # TODO this is where retry of an entire step would happen
    614 if inflight:
    615     # if we got here means we timed out

File ~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/thread.py:58, in _WorkItem.run(self)
     55     return
     57 try:
---> 58     result = self.fn(*self.args, **self.kwargs)
     59 except BaseException as exc:
     60     self.future.set_exception(exc)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:3596, in RunnableBindingBase.invoke(self, input, config, **kwargs)
   3590 def invoke(
   3591     self,
   3592     input: Input,
   3593     config: Optional[RunnableConfig] = None,
   3594     **kwargs: Optional[Any],
   3595 ) -> Output:
-> 3596     return self.bound.invoke(
   3597         input,
   3598         self._merge_configs(config),
   3599         **{**self.kwargs, **kwargs},
   3600     )

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1774, in RunnableSequence.invoke(self, input, config)
   1772 try:
   1773     for i, step in enumerate(self.steps):
-> 1774         input = step.invoke(
   1775             input,
   1776             # mark each step as a child run
   1777             patch_config(
   1778                 config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
   1779             ),
   1780         )
   1781 # finish the root run
   1782 except BaseException as e:

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/passthrough.py:415, in RunnableAssign.invoke(self, input, config, **kwargs)
    409 def invoke(
    410     self,
    411     input: Dict[str, Any],
    412     config: Optional[RunnableConfig] = None,
    413     **kwargs: Any,
    414 ) -> Dict[str, Any]:
--> 415     return self._call_with_config(self._invoke, input, config, **kwargs)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:975, in Runnable._call_with_config(self, func, input, config, run_type, **kwargs)
    971     context = copy_context()
    972     context.run(var_child_runnable_config.set, child_config)
    973     output = cast(
    974         Output,
--> 975         context.run(
    976             call_func_with_variable_args,
    977             func,  # type: ignore[arg-type]
    978             input,  # type: ignore[arg-type]
    979             config,
    980             run_manager,
    981             **kwargs,
    982         ),
    983     )
    984 except BaseException as e:
    985     run_manager.on_chain_error(e)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/config.py:323, in call_func_with_variable_args(func, input, config, run_manager, **kwargs)
    321 if run_manager is not None and accepts_run_manager(func):
    322     kwargs["run_manager"] = run_manager
--> 323 return func(input, **kwargs)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/passthrough.py:402, in RunnableAssign._invoke(self, input, run_manager, config, **kwargs)
    389 def _invoke(
    390     self,
    391     input: Dict[str, Any],
   (...)
    394     **kwargs: Any,
    395 ) -> Dict[str, Any]:
    396     assert isinstance(
    397         input, dict
    398     ), "The input to RunnablePassthrough.assign() must be a dict."
    400     return {
    401         **input,
--> 402         **self.mapper.invoke(
    403             input,
    404             patch_config(config, callbacks=run_manager.get_child()),
    405             **kwargs,
    406         ),
    407     }

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:2339, in RunnableParallel.invoke(self, input, config)
   2326     with get_executor_for_config(config) as executor:
   2327         futures = [
   2328             executor.submit(
   2329                 step.invoke,
   (...)
   2337             for key, step in steps.items()
   2338         ]
-> 2339         output = {key: future.result() for key, future in zip(steps, futures)}
   2340 # finish the root run
   2341 except BaseException as e:

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:2339, in <dictcomp>(.0)
   2326     with get_executor_for_config(config) as executor:
   2327         futures = [
   2328             executor.submit(
   2329                 step.invoke,
   (...)
   2337             for key, step in steps.items()
   2338         ]
-> 2339         output = {key: future.result() for key, future in zip(steps, futures)}
   2340 # finish the root run
   2341 except BaseException as e:

File ~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/_base.py:456, in Future.result(self, timeout)
    454     raise CancelledError()
    455 elif self._state == FINISHED:
--> 456     return self.__get_result()
    457 else:
    458     raise TimeoutError()

File ~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/_base.py:401, in Future.__get_result(self)
    399 if self._exception:
    400     try:
--> 401         raise self._exception
    402     finally:
    403         # Break a reference cycle with the exception in self._exception
    404         self = None

File ~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/thread.py:58, in _WorkItem.run(self)
     55     return
     57 try:
---> 58     result = self.fn(*self.args, **self.kwargs)
     59 except BaseException as exc:
     60     self.future.set_exception(exc)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:1774, in RunnableSequence.invoke(self, input, config)
   1772 try:
   1773     for i, step in enumerate(self.steps):
-> 1774         input = step.invoke(
   1775             input,
   1776             # mark each step as a child run
   1777             patch_config(
   1778                 config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
   1779             ),
   1780         )
   1781 # finish the root run
   1782 except BaseException as e:

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:167, in BaseOutputParser.invoke(self, input, config)
    163 def invoke(
    164     self, input: Union[str, BaseMessage], config: Optional[RunnableConfig] = None
    165 ) -> T:
    166     if isinstance(input, BaseMessage):
--> 167         return self._call_with_config(
    168             lambda inner_input: self.parse_result(
    169                 [ChatGeneration(message=inner_input)]
    170             ),
    171             input,
    172             config,
    173             run_type="parser",
    174         )
    175     else:
    176         return self._call_with_config(
    177             lambda inner_input: self.parse_result([Generation(text=inner_input)]),
    178             input,
    179             config,
    180             run_type="parser",
    181         )

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/base.py:975, in Runnable._call_with_config(self, func, input, config, run_type, **kwargs)
    971     context = copy_context()
    972     context.run(var_child_runnable_config.set, child_config)
    973     output = cast(
    974         Output,
--> 975         context.run(
    976             call_func_with_variable_args,
    977             func,  # type: ignore[arg-type]
    978             input,  # type: ignore[arg-type]
    979             config,
    980             run_manager,
    981             **kwargs,
    982         ),
    983     )
    984 except BaseException as e:
    985     run_manager.on_chain_error(e)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/runnables/config.py:323, in call_func_with_variable_args(func, input, config, run_manager, **kwargs)
    321 if run_manager is not None and accepts_run_manager(func):
    322     kwargs["run_manager"] = run_manager
--> 323 return func(input, **kwargs)

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:168, in BaseOutputParser.invoke.<locals>.<lambda>(inner_input)
    163 def invoke(
    164     self, input: Union[str, BaseMessage], config: Optional[RunnableConfig] = None
    165 ) -> T:
    166     if isinstance(input, BaseMessage):
    167         return self._call_with_config(
--> 168             lambda inner_input: self.parse_result(
    169                 [ChatGeneration(message=inner_input)]
    170             ),
    171             input,
    172             config,
    173             run_type="parser",
    174         )
    175     else:
    176         return self._call_with_config(
    177             lambda inner_input: self.parse_result([Generation(text=inner_input)]),
    178             input,
    179             config,
    180             run_type="parser",
    181         )

File ~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/output_parsers/base.py:219, in BaseOutputParser.parse_result(self, result, partial)
    206 def parse_result(self, result: List[Generation], *, partial: bool = False) -> T:
    207     """Parse a list of candidate model Generations into a specific format.
    208 
    209     The return value is parsed from only the first Generation in the result, which
   (...)
    217         Structured output.
    218     """
--> 219     return self.parse(result[0].text)

Cell In[1], line 88, in LLMCompilerPlanParser.parse(self, text)
     86 parser = ActionParserFSM()
     87 graph_dict = {}
---> 88 for task in parser.parse(text):
     89     idx = int(task["task_index"])
     91     task = instantiate_task(
     92         tools=self.tools,
     93         idx=idx,
     94         tool_name=task["tool_name"],
     95         args=task["args"],
     96     )

Cell In[1], line 29, in ActionParserFSM.parse(self, text)
     27 def parse(self, text: str):
     28     for char in text:
---> 29         action = self.process_char(char)
     30         if action:
     31             yield action

Cell In[1], line 61, in ActionParserFSM.process_char(self, char)
     59 elif self.state == "COMMENT":
     60     if char == "\n":
---> 61         action = self.save_action()
     62         self.reset()
     63     else:

Cell In[1], line 69, in ActionParserFSM.save_action(self)
     67 def save_action(self):
     68     if self.task_index and self.action:
---> 69         parsed_action = json.loads(self.action.strip())
     70         tool_name, args = next(iter(parsed_action.items()))
     71         return {
     72             "task_index": int(self.task_index),
     73             "tool_name": tool_name,
     74             "args": args,
     75         }

File ~/.pyenv/versions/3.11.2/lib/python3.11/json/__init__.py:346, in loads(s, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)
    341     s = s.decode(detect_encoding(s), 'surrogatepass')
    343 if (cls is None and object_hook is None and
    344         parse_int is None and parse_float is None and
    345         parse_constant is None and object_pairs_hook is None and not kw):
--> 346     return _default_decoder.decode(s)
    347 if cls is None:
    348     cls = JSONDecoder

File ~/.pyenv/versions/3.11.2/lib/python3.11/json/decoder.py:337, in JSONDecoder.decode(self, s, _w)
    332 def decode(self, s, _w=WHITESPACE.match):
    333     """Return the Python representation of ``s`` (a ``str`` instance
    334     containing a JSON document).
    335 
    336     """
--> 337     obj, end = self.raw_decode(s, idx=_w(s, 0).end())
    338     end = _w(s, end).end()
    339     if end != len(s):

File ~/.pyenv/versions/3.11.2/lib/python3.11/json/decoder.py:355, in JSONDecoder.raw_decode(self, s, idx)
    353     obj, end = self.scan_once(s, idx)
    354 except StopIteration as err:
--> 355     raise JSONDecodeError("Expecting value", s, err.value) from None
    356 return obj, end

JSONDecodeError: Expecting value: line 1 column 1 (char 0)

In [ ]: