通过构建有向无环图DAG来表示任务之间的依赖关系，LLM Compiler能够实现任务的并行执行，从而大幅降低总执行时间。本文将详细介绍LLM Compiler的原理、实现过程以及其在实际应用中的优势。

在上篇文章《AI大模型实战篇：AI Agent设计模式 – Plan & Execute》中，风叔结合原理和具体源代码，详细介绍了AI Agent设计模式中的Plan-and-Execute。但是Plan-and-execute的局限性在于，每个任务是按顺序执行的，这可能会导致总执行时间的增加。

一种有效改进的办法是将每个任务表示为有向无环图DAG，这样可以让多个任务并行执行，大幅降低执行总时间。

这就是本篇文章风叔将为大家介绍的AI Agent设计模式，LLM Compiler。

01 LLM Compiler的概念

LLM Compiler是伯克利大学的SqueezeAILab于2023年12月提出的新项目。这个项目在ReWOO引入的变量分配的基础上，进一步训练大语言模型生成一个有向无环图（Directed Acyclic Graph，DAG，如下图所示）类的规划。DAG可以明确各步骤任务之间的依赖关系，从而并行执行任务，实现类似处理器“乱序执行”的效果，可以大幅加速AI Agent完成任务的速度。

1. LLM Compiler设计模式主要有以下组件：

Planner：输出流式传输任务的DAG，每个任务都包含一个工具、参数和依赖项列表。相比ReWOO的Planner，依赖项列表是最大的不同。
Task Fetching Unit：调度并执行任务，一旦满足任务的依赖性，该单元就会安排任务。由于许多工具涉及对搜索引擎或LLM的其他调用，因此额外的并行性可以显著提高速度。
Joiner：由LLM根据整个历史记录（包括任务执行结果），决定是否响应最终答案或是否将进度重新传递回Planner。

2. 下图是LLM Compiler的原理：

Planner接收来自用户的输入，输出流式传输任务的DAG
Task Fetching Unit从式传输任务DAG中读取任务，通过处理工具并行执行
Task Fetching Unit将状态和结果传递给Joiner（或Replanner），Joiner来决定是将结果输出给用户，还是增加更多任务交由Task Fetching Unit处理

02 LLM Compiler的实现过程

下面，风叔通过实际的源码，详细介绍LLM Compiler模式的实现方法。

第一步构建工具Tools

首先，我们要定义Agent需要使用的工具。在这个例子中，我们将使用搜索引擎 + 计算器这两个工具。

_get_pass(“TAVILY_API_KEY”)

第二步构建Planner

Planner接收用户输入，并生成一个待执行的任务清单的DAG。

以下代码构建了Planner的提示模板，并将其与 LLM 和输出解析器组合在一起，输出解析器处理以下形式的任务列表。在Planner中，我们同时定义了replanner的Prompt，这个prompt提出了三项核心的约束

启动当前计划时，应该从概述下一个计划策略的“Thought”开始
在当前计划中，绝不应该重复上一个计划中已经执行的操作
必须从上一个任务索引的末尾继续任务索引，不要重复任务索引

def create_planner(
llm: BaseChatModel, tools: Sequence[BaseTool], base_prompt: ChatPromptTemplate
):
tool_descriptions = “n”.join(
f”{i+1}. {tool.description}n”
for i, tool in enumerate(
tools
) # +1 to offset the 0 starting index, we want it count normally from 1.
)
planner_prompt = base_prompt.partial(
replan=””,
num_tools=len(tools)
+ 1, # Add one because we’re adding the join() tool at the end.
tool_descriptions=tool_descriptions,
)
replanner_prompt = base_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.”,
num_tools=len(tools) + 1,
tool_descriptions=tool_descriptions,
)

def should_replan(state: list):
# Context is passed as a system message
return isinstance(state[-1], SystemMessage)

def wrap_messages(state: list):
return {“messages”: state}

def wrap_and_get_last_index(state: list):
next_task = 0
for message in state[::-1]:
if isinstance(message, FunctionMessage):
next_task = message.additional_kwargs[“idx”] + 1
break
state[-1].content = state[-1].content + f” – Begin counting at : {next_task}”
return {“messages”: state}

这个部分负责安排任务，它接收以下格式的数据流。

tool:BaseTool,

dependencies:number[]

其核心思想是，一旦满足依赖关系，就开始执行工具，可以通过多线程实现。下面这段代码的关键就在于schedule_tasks，会将所有任务处理成有向无环图。在当前任务存在尚未完成的依赖关系时，放入pending task；在当前任务所有依赖关系都已完成时，执行任务。

@as_runnable

def schedule_task(task_inputs, config):
task: Task = task_inputs[“task”]
observations: Dict[int, Any] = task_inputs[“observations”]
try:
observation = _execute_task(task, observations, config)
except Exception:
import traceback

observations[task[“idx”]] = observation
def schedule_pending_task(task: Task, observations: Dict[int, Any], retry_after: float = 0.2):
while True:
deps = task[“dependencies”]
if deps and (any([dep not in observations for dep in deps])):
# Dependencies not yet satisfied
time.sleep(retry_after)
continue
schedule_task.invoke({“task”: task, “observations”: observations})
break

@as_runnable
def schedule_tasks(scheduler_input: SchedulerInput) -> List[FunctionMessage]:
“””Group the tasks into a DAG schedule.”””
tasks = scheduler_input[“tasks”]
args_for_tasks = {}
messages = scheduler_input[“messages”]

observations = _get_observations(messages)
task_names = {}
originals = set(observations)

futures = []
retry_after = 0.25 # Retry every quarter second
with ThreadPoolExecutor() as executor:
for task in tasks:
deps = task[“dependencies”]
task_names[task[“idx”]] = (
task[“tool”] if isinstance(task[“tool”], str) else task[“tool”].name
)
args_for_tasks[task[“idx”]] = task[“args”]
if (
# Depends on other tasks
deps
and (any([dep not in observations for dep in deps]))
):
futures.append(
executor.submit(
schedule_pending_task, task, observations, retry_after
)
)
else:
# No deps or all deps satisfied，can schedule now
schedule_task.invoke(dict(task=task, observations=observations))
# futures.append(executor.submit(schedule_task.invoke dict(task=task, observations=observations)))

# All tasks have been submitted or enqueued
# Wait for them to complete
wait(futures)
# Convert observations to new tool messages to add to the state
new_observations = {
k: (task_names[k], args_for_tasks[k], observations[k])
for k in sorted(observations.keys() – originals)
}

tool_messages = [
FunctionMessage(
name=name, content=str(obs), additional_kwargs={“idx”: k, “args”: task_args}
)
for k, (name, task_args, obs) in new_observations.items()
]

import itertools

@as_runnable
def plan_and_schedule(state):
messages = state[“messages”]
tasks = planner.stream(messages)
# Begin executing the planner immediately
try:
tasks = itertools.chain([next(tasks)], tasks)
except StopIteration:
# Handle the case where tasks is empty.
tasks = iter([])
scheduled_tasks = schedule_tasks.invoke(
{
“messages”: messages,
“tasks”: tasks,
}
)
return {“messages”: [scheduled_tasks]}

第四步构建Joiner

前面我们构建了Planner和Task Fetching Unit，下一步我们需要构建Joiner来处理工具的输出，以及决定是否需要使用新的计划并开启新的循环。

class FinalResponse(BaseModel):
“””The final response/answer.”””
response: str

class Replan(BaseModel):
feedback: str = Field(
description=”Analysis of the previous attempts and recommendations on what needs to be fixed.”
)

class JoinOutputs(BaseModel):
“””Decide whether to replan or whether you can return the final response.”””
thought: str = Field(
description=”The chain of thought reasoning for the selected action”
)
action: Union[FinalResponse, Replan]

joiner_prompt = hub.pull(“wfh/llm-compiler-joiner”).partial(
examples=””
) # You can optionally add examples

llm = ChatOpenAI(model=”gpt-4-turbo-preview”)
runnable = create_structured_output_runnable(JoinOutputs, llm, joiner_prompt)

如果Agent需要继续循环，我们需要选择状态机内的最新消息，并按照Planner的要求输出相应的格式。

return {“messages”: selected[::-1]}
joiner = select_recent_messages | runnable | _parse_joiner_output
input_messages = [HumanMessage(content=example_question)] + tool_messages
joiner.invoke(input_messages)

下面，我们构建流程图，将Planner、Task Fetching Unit、Joiner等节点添加进来，循环执行并输出结果。

class State(TypedDict):
messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

graph_builder.add_node(“plan_and_schedule”, plan_and_schedule)
graph_builder.add_node(“join”, joiner)
graph_builder.add_edge(“plan_and_schedule”, “join”)

def should_continue(state):
messages = state[“messages”]
if isinstance(messages[-1], AIMessage):
return END
return “plan_and_schedule”

graph_builder.add_conditional_edges(
start_key=”join”,
# Next, we pass in the function that will determine which node is called next.
condition=should_continue,
)

graph_builder.add_edge(START, “plan_and_schedule”)
chain = graph_builder.compile()

总结

通过前面三篇文章，按照递进关系，风叔依次介绍了REWOO、Plan-and-Execute和LLM Compiler三种更侧重规划能力的AI Agent设计模式。从最初的ReAct模式出发，加入规划能力即演变成REWOO；再加上Replan能力即演变成Plan-and-Execute；最后再加上DAG和并行处理能力，即演变成LLM Compiler。

题图来自Unsplash，基于 CC0 协议。

AI大模型实战篇：AI Agent设计模式 – LLM Compiler，Agent的并行处理器