LATS,Language Agent Tree Search,语言智能体树搜索(ICML 2024 论文)
用 LangGraph 搭建一棵思考树,每一步都生成多个备选方案,像 DFS 一样探索所有可能路径;让大模型当裁判投票,挑最好的那条思考路径,最后用这条最优路径输出答案
不像普通 ReAct 一条路走到黑,ToT 是把思考拆成多分支树,每一步生成多个候选方案,像深度优先搜索一样探索多条路径,再投票选出最好的一条路线
superfluoussuːˈpɜːfluəsadj. 多余的,过剩的;奢侈的
外部参数
max_candidates:重规划下一个 step 时,只需要前 x 个候选 step
max_paths:最大路径,超过最大路径后,只能跳 vote 决策;否则会回到 run 节点,继续生成当前步骤答案->决策出下一步候选步骤列表
planning agent
frompydanticimportBaseModel,Fieldfromlangchain_core.promptsimportChatPromptTemplatefromlangchain_google_genaiimportChatGoogleGenerativeAIclassPlan(BaseModel):"""Plan to follow in future"""steps:list[str]=Field(description="different steps to follow, should be in sorted order")system_prompt_template=(# come up with a step by step plan,制定一个分步计划# final step 的 result 应该是最终 answer"For the given task, come up with a step by step plan.\n""This plan should involve individual tasks, that if executed correctly will ""yield the correct answer. Do not add any superfluous steps.\n""The result of the final step should be the final answer. Make sure that each ""step has all the information needed - do not skip steps.")planner_prompt=ChatPromptTemplate.from_messages([("system",system_prompt_template),("user","Prepare a plan how to solve the following task:\n{task}\n")])planner=planner_prompt|ChatGoogleGenerativeAI(model="gemini-2.5-flash",temperature=1.0).with_structured_output(Plan)构建 agent 用于执行 plan 中的一个 step
fromlangchain_community.agent_toolkits.load_toolsimportload_toolsfromlanggraph.prebuiltimportcreate_react_agent llm=ChatGoogleGenerativeAI(model="gemini-2.5-flash")# DDG = DuckDuckGo,`ddg-search` 就是 DuckDuckGo 网页搜索工具(MCP 工具)# DuckDuckGo 是主打隐私保护的英文搜索引擎,官网:[duckduckgo.com](https://duckduckgo.com)# 适合 Agent 回答时效性问题(新闻、文档、产品信息),弥补大模型知识截止的缺陷tools=load_tools(tool_names=["ddg-search","arxiv","wikipedia"],llm=llm)system_prompt=(# 给定一个用于解决一个 task 的通用 plan,和一个特定 step,work on this step,进行这一步# 使用 tools 来二确自己"You're a smart assistant that carefully helps to solve complex tasks.\n"" Given a general plan to solve a task and a specific step, work on this step. "" Don't assume anything, keep in minds things might change and always try to ""use tools to double-check yourself.\nUse Search to gather ""information about common facts, fresh events and news, use Arxiv to get ""ideas on recent research and use Wikipedia for common knowledge.")step_template=("Given the task and the plan, try to execute on a specific step of the plan.\n""TASK:\n{task}\n\nPLAN:\n{previous_steps}\n\nSTEP TO EXECUTE:\n{step}\n")prompt_template=ChatPromptTemplate.from_messages([("system",system_prompt),("user",step_template),])execution_agent=prompt_template|create_react_agent(model=llm,tools=tools)定义一个单独的结构来跟踪将要探索的选项树
fromtypingimportOptionalclassTreeNode:def__init__(self,node_id:int,# step:计划 / 动作 / 思考(要做什么)# step 是 当前步骤的名字# step_output:观测、返回结果(做完之后拿到了什么)step:str,step_output:Optional[str]=None,parent:Optional["TreeNode"]=None,):self.node_id=node_id self.step=step self.step_output=step_output self.parent=parent self.children=[]self.final_response=Nonedef__repr__(self):parent_id=self.parent.node_idifself.parentelse"None"returnf"Node_id:{self.node_id}, parent:{parent_id},{len(self.children)}children."defget_full_plan(self)->str:"""Returns formatted plan with step numbers and past results."""steps=[]node=selfwhilenode.parent:steps.append((node.step,node.step_output))node=node.parent full_plan=[]fori,(step,result)inenumerate(steps[::-1]):ifresult:full_plan.append(f"#{i+1}. Planned step:{step}\nResult:{result}\n")return"\n".join(full_plan)定义额外的状态
probableˈprɒbəb(ə)l
adj. 很可能发生(或存在等)的,大概的
n. 可能获胜的人(或动物),可能入选者
importoperatorfromcollectionsimportdequefromtypingimportAnnotated,TypedDictclassPlanEvaluation(BaseModel):""""""is_final:bool=Field(description="Whether the task is solved or not. Plan is not final if additional steps are required.",default=False)score:float=Field(description="A score from 0. to 1.0 that evaluates the quality of the plan and how probable it is to solve the task with this plan.",default=False)classPlanState(TypedDict):task:strroot:TreeNode queue:deque[TreeNode]current_node:TreeNode next_node:TreeNode is_current_node_final:boolpaths_explored:Annotated[int,operator.add]visited_ids:set[int]max_id:intcandidates:Annotated[list[str],operator.add]best_candidate:str创建一个replanning agent,考虑之前的steps和适当调节plan
create a replanning agent that takes into account the previous steps and adjusts the plan accordingly:
fromlangchain_core.promptsimportChatPromptTemplate# 作用:在执行完一步之后,重新规划,输出多个候选下一步动作# 不是只输出 1 个方案!会生成多个可选路径,这就是树的分支classReplanStep(BaseModel):"""Replanned next step in the plan."""steps:list[str]=Field(description="different options of the proposed next step")llm_replanner=llm.with_structured_output(ReplanStep)replanner_prompt_template=("Suggest next action in the plan. Do not add any superfluous steps.\n""If you think no actions are needed, just return an empty list of steps. ""TASK: {task}\n PREVIOUS STEPS WITH OUTPUTS: {current_plan}")replanner_prompt=ChatPromptTemplate.from_messages([("system","You're a helpful assistant. You goal is to help with planning actions to solve the task. Do not solve the task itself."),("user",replanner_prompt_template)])replanner=replanner_prompt|llm_replannerfinal step,投票选出最好的 candidate
We also need a final step that includes a major voting for the best candidate:
fromlangchain_core.promptsimportPromptTemplatefromlangchain_core.output_parsersimportStrOutputParser prompt_voting=PromptTemplate.from_template("Pick the best solution for a given task. ""\nTASK:{task}\n\nSOLUTIONS:\n{candidates}\n")# 模型投票选出最优分支# 输入:多条不同路径的候选方案# LLM 充当评委,从多个候选里挑选最优的下一步方案# 结构化输出,返回选中的最佳候选编号# ToT 的核心亮点:同一个步骤多条路线并行探索,用 LLM 做评判、剪枝,抛弃差的分支,而不是一条路硬走def_vote_for_the_best_option(state):candidates=state.get("candidates",[])ifnotcandidates:return{"best_response":None}all_candidates=[]fori,candidateinenumerate(candidates):all_candidates.append(f"OPTION{i+1}:{candidate}")response_schema={"type":"STRING",# 枚举,大模型只能从这些数字中选"enum":[str(i+1)foriinrange(len(all_candidates))],}llm_enum=ChatGoogleGenerativeAI(model="gemini-2.5-flash",).with_structured_output(response_schema)result=(prompt_voting|llm_enum).invoke({"candidates":"\n".join(all_candidates),"task":state["task"]})return{"best_candidate":candidates[int(result)-1]}_vote_for_the_best_option({"candidates":["1","5","4"],"task":"How much is 2+2?"})汇总上述逻辑,构建 agent 图
fromtypingimportLiteralfromlanggraph.graphimportStateGraph,START,ENDfromlangchain_core.runnablesimportRunnableConfigfromlangchain_core.output_parsersimportStrOutputParserfromlanggraph.typesimportCommand# 不是当前单一步骤的答案;是把【整条路径全部步骤 + 每一步执行结果】一次性交给 LLM,生成针对原始 task 的最终完整答案final_prompt=PromptTemplate.from_template("You're a helpful assistant that has executed on a plan.""Given the results of the execution, prepare the final response.\n""Don't assume anything\nTASK:\n{task}\n\nPLAN WITH RESUlTS:\n{plan}\n""FINAL RESPONSE:\n")responder=final_prompt|llm|StrOutputParser()# 这是一条单链(链式树,也就是链表结构),不是多分支树,一条从头到尾串联起来的父子链条# queue:A -> B -> Casyncdef_build_initial_plan(state:PlanState)->PlanState:plan=awaitplanner.ainvoke(state["task"])queue=deque()root=TreeNode(step=plan.steps[0],node_id=1)queue.append(root)current_root=rootfori,stepinenumerate(plan.steps[1:]):child=TreeNode(node_id=i+2,step=step,parent=current_root)current_root.children.append(child)queue.append(child)current_root=child# `max_id` 记录当前已经使用过最大的 node_id。后面新增子节点时,拿这个值 +1,保证每个 TreeNode 的`node_id`全局唯一,不会重复return{"root":root,"queue":queue,"max_id":i+2}asyncdef_run_node(state:PlanState,config:RunnableConfig):# 用 next_node 显式指定下一个要执行的节点node=state.get("next_node")visited_ids=state.get("visited_ids",set())queue=state["queue"]ifnodeisNone:whilequeueandnotnode:node=state["queue"].popleft()ifnode.node_idinvisited_ids:node=Noneifnotnode:returnCommand(goto="vote",update={})step=awaitexecution_agent.ainvoke({"previous_steps":node.get_full_plan(),"step":node.step,"task":state["task"]})node.step_output=step["messages"][-1].content visited_ids.add(node.node_id)return{"current_node":node,"queue":queue,"visited_ids":visited_ids,"next_node":None}asyncdef_plan_next(state:PlanState,config:RunnableConfig)->PlanState:# max_candidates:重规划时,只需要前 x 个候选max_candidates=config["configurable"].get("max_candidates",1)node=state["current_node"]next_step=awaitreplanner.ainvoke({"task":state["task"],"current_plan":node.get_full_plan()})ifnotnext_step.steps:return{"is_current_node_final":True}# `max_id` 记录当前已经使用过最大的 node_id。后面新增子节点时,拿这个值 +1,保证每个 TreeNode 的`node_id`全局唯一,不会重复max_id=state["max_id"]forstepinnext_step.steps[:max_candidates]:child=TreeNode(node_id=max_id+1,step=step,parent=node)max_id+=1node.children.append(child)state["queue"].append(child)return{"is_current_node_final":False,"next_node":child,"max_id":max_id}# candidates:存放本次探索路径得到的候选答案,注意,是最终交可能会选定给到用户的答案asyncdef_get_final_response(state:PlanState)->PlanState:node=state["current_node"]final_response=awaitresponder.ainvoke({"task":state["task"],"plan":node.get_full_plan()})node.final_response=final_responsereturn{"paths_explored":1,"candidates":[final_response]}def_should_create_final_response(state:PlanState)->Literal["run","generate_response"]:return"generate_response"ifstate["is_current_node_final"]else"run"def_should_continue(state:PlanState,config:RunnableConfig)->Literal["run","vote"]:max_paths=config["configurable"].get("max_paths",30)# max_paths 不能太深ifstate.get("paths_explored",1)>=max_paths:return"vote"ifstate["queue"]orstate.get("next_node"):return"run"return"vote"builder=StateGraph(PlanState)builder.add_node("initial_plan",_build_initial_plan)builder.add_node("run",_run_node)builder.add_node("plan_next",_plan_next)builder.add_node("generate_response",_get_final_response)builder.add_node("vote",_vote_for_the_best_option)builder.add_edge(START,"initial_plan")builder.add_edge("initial_plan","run")builder.add_edge("run","plan_next")# _should_create_final_response 跳转 run 或者 generate_responsebuilder.add_conditional_edges("plan_next",_should_create_final_response)# _should_continue 跳转 run 或者 votebuilder.add_conditional_edges("generate_response",_should_continue)builder.add_edge("vote",END)graph=builder.compile()fromIPython.displayimportImage,display display(Image(graph.get_graph().draw_mermaid_png()))