LangGraph 和 ToT，Tree of Thoughts 同一步骤多条路线并行探索，用LLM做评判和减枝

LATS，Language Agent Tree Search，语言智能体树搜索（ICML 2024 论文）

用 LangGraph 搭建一棵思考树，每一步都生成多个备选方案，像 DFS 一样探索所有可能路径；让大模型当裁判投票，挑最好的那条思考路径，最后用这条最优路径输出答案

不像普通 ReAct 一条路走到黑，ToT 是把思考拆成多分支树，每一步生成多个候选方案，像深度优先搜索一样探索多条路径，再投票选出最好的一条路线

from typing import TypedDict, Optional, Annotated
from collections import deque
import operator
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field
from langgraph.graph import StateGraph, START, END
from langgraph.runnables import RunnableConfig

# ===================== 1. 定义重规划输出结构 =====================
# 作用：在执行完一步之后，重新规划，输出多个候选下一步动作
# 不是只输出 1 个方案！会生成多个可选路径，这就是树的分支
class ReplanStep(BaseModel):
    steps: list[str] = Field(description="different options of the proposed next step")

# 重规划Prompt：执行完一步，生成多个候选下一步
replanner_prompt_template = ChatPromptTemplate.from_messages([
    ("system", """You're a helpful assistant. Your goal is to help with planning actions to solve the task. Do not solve the task itself."""),
    ("user", """TASK: {task}
PREVIOUS STEPS WITH OUTPUTS: {current_plan}
Suggest next action in the plan. Do not add any superfluous steps.
If you think no actions are needed, just return an empty list of steps."""),
])

# ===================== 2. TreeNode 树节点，保存搜索树 =====================
# 用来保存搜索树状态：每一个节点代表「某一步规划 + 执行结果」
# `get_full_plan()`：从当前节点向上回溯父节点，拿到从根到当前节点完整历史路径
class TreeNode:
    def __init__(
        self,
        node_id: int,
        # step 是 当前节点的步骤的名字
        step: str,
        step_output: Optional[str] = None,
        parent: Optional["TreeNode"] = None,
    ):
        self.node_id = node_id
  	    # step：计划 / 动作 / 思考（要做什么）
        # step_output：观测、返回结果（做完之后拿到了什么）
        self.step = step
        self.step_output = step_output
        self.parent = parent
        self.children: list[TreeNode] = []
        self.final_response = None

    # 方便打印调试，查看整条思考链路
    def __repr__(self):
        parent_id = self.parent.node_id if self.parent else "None"
        return f"Node(id: {self.node_id}, parent: {parent_id}, {len(self.children)} children)"

    def get_full_plan(self) -> str:
        """从当前节点向上回溯，拿到从根到当前节点完整路径"""
        full_plan = []
        node = self
        while node.parent:
            full_plan.append((node.step, node.step_output))
            node = node.parent
        full_plan.reverse()
        return "\n".join(f"#{i+1}. Planned step: {step}\nResult: {result}"
                         for i, (step, result) in enumerate(full_plan))

# ===================== 3. 定义LangGraph状态PlanState =====================
class PlanState(TypedDict):
    task: str
    root: TreeNode
    queue: deque[TreeNode]          # DFS待探索节点队列
    visited_ids: set[int]           # 已访问节点，防重复
    paths_explored: Annotated[list[str], operator.add] # 记录所有探索路径，追加而非覆盖
    candidates: Annotated[list[int], operator.add]
    max_candidates: int
    max_paths: int
    next_node: Optional[TreeNode]

# ===================== 4. 图节点函数：_run_node 执行+生成子候选 =====================
# 深度优先搜索 DFS
# 1. 从队列取出待探索节点
# 2. 调用`replanner`，生成多个候选下一步（多个子节点）
# 3. 把所有候选子节点加入树，放进待探索队列
# 4. 标记节点已访问，防止重复探索
async def _run_node(state: PlanState, config: RunnableConfig):
    queue = state["queue"]
    if not queue:
        return {"next_node": None}
    node = queue.popleft()
    if node.node_id in state["visited_ids"]:
        return {"next_node": None}
    state["visited_ids"].add(node.node_id)

    # 调用replanner，生成多个下一步候选动作
    llm_replanner = replanner_prompt_template | llm.with_structured_output(ReplanStep)
    current_plan = node.get_full_plan()
    replan_result = await llm_replanner.ainvoke({
        "task": state["task"],
        "current_plan": current_plan
    })
    next_steps = replan_result.steps

    # 生成子节点，挂到树上，放进队列
    max_id = state["max_candidates"]
    child_nodes = []
    for step in next_steps:
        max_id += 1
        child = TreeNode(node_id=max_id, step=step, parent=node)
        node.children.append(child)
        child_nodes.append(child)
    
    # 深度优先：新节点塞到队列头部
    for child in reversed(child_nodes):
        queue.appendleft(child)

    return {
        "queue": queue,
        "visited_ids": state["visited_ids"],
        "max_candidates": max_id,
        "next_node": node
    }

# ===================== 5. 投票节点：vote_for_the_best 大模型选最优候选 =====================
# 模型投票选出最优分支
# 输入：多条不同路径的候选方案
# LLM 充当评委，从多个候选里挑选最优的下一步方案
# 结构化输出，返回选中的最佳候选编号
# ToT 的核心亮点：同一个步骤多条路线并行探索，用 LLM 做评判、剪枝，抛弃差的分支，而不是一条路硬走
async def vote_for_the_best(state: PlanState):
    candidates = state.get("candidates", [])
    if not candidates:
        return {"best_response": None}
    prompt_voting = ChatPromptTemplate.from_template("""
Pick the best solution for a given task.
TASK:{task}
{nSOLUTIONS:\n{candidates}}
""")
    # schema 锁死输出，避免模型输出冗长解释，方便后面代码直接取编号索引
    response_schema = {
        "type": "STRING",
        # `enum` = 枚举，限定模型只能输出列表里面预先写好的值，不允许输出别的文字
        "enum": [str(i+1) for i in range(len(candidates))]
    }
   
    llm_enum = ChatVertexAI(model_name="gemini-2.0-flash-001")
    chain = prompt_voting | llm_enum | StrOutputParser()
    
    result = await chain.ainvoke({
        "task": state["task"],
        "candidates": "\n".join(candidates)
    })
    return {"best_candidate": result}

# ===================== 6. 初始化节点：build_initial_plan =====================
async def build_initial_plan(state: PlanState):
    planner_prompt = ChatPromptTemplate.from_template("""
Create an initial plan to solve the following task.
Only output your plan steps, do NOT execute or give final answer.
TASK: {task}
""")
    chain = planner_prompt | llm | StrOutputParser()
    root_plan = await chain.ainvoke({"task": state["task"]})
    root = TreeNode(node_id=1, step=root_plan)
    queue = deque()
    queue.append(root)
    return {
        "root": root,
        "queue": queue,
        "visited_ids": set(),
        "max_candidates": 1,
    }

# ===================== 7. 终止判断 + 最终生成响应 =====================
def should_continue(state: PlanState):
    if state["queue"] or state.get("next_node"):
        return "run"
    else:
        return "vote"

# 前面所有树搜索、多分支探索、LLM 投票，都是找最优路径；这个函数拿投票选出的最优规划路径，交给大模型，生成给用户看的最终回答
# 构造提示词，告诉 LLM：已经探索了很多方案，基于这个选定好的 plan，写最终答案
async def generate_final_response(state: PlanState):
    final_prompt = ChatPromptTemplate.from_template("""
You've explored multiple solutions. Pick the best one and write the final answer.
Task:{task}, plan:{plan}
""")
    chain = final_prompt | llm | StrOutputParser()
    
    # state["paths_explored"] 是个字符串列表
    best_plan = state["paths_explored"][int(state["best_candidate"])-1]
    resp = await chain.ainvoke({"task": state["task"], "plan": best_plan})
    return {"final_response": resp}

# ===================== 8. 构建LangGraph图 =====================
builder = StateGraph(PlanState)
builder.add_node("initial_plan", build_initial_plan)
builder.add_node("run", _run_node)
builder.add_node("vote", vote_for_the_best)
builder.add_node("generate_response", generate_final_response)

builder.add_edge(START, "initial_plan")
builder.add_edge("initial_plan", "run")
builder.add_conditional_edges("run", should_continue, {"run": "run", "vote": "vote"})
builder.add_edge("vote", "generate_response")
builder.add_edge("generate_response", END)

graph = builder.compile()

开始 → 生成初始规划 initial_plan
    ↓
执行当前步骤 run_node
    ↓
重新规划 replanner：生成多个下一步候选方案（多个分支）
    ↓
把候选放进搜索树TreeNode，用队列做深度优先遍历
    ↓
投票vote_for_the_best：大模型给多个方案打分，选出最优分支
    ↓
判断是否结束 should_continue：
    ✅ 没结束：回到执行节点，继续往下搜索树
    ✅ 结束：生成最终回答 generate_final_response

should_continue：判断任务是否完成，是否停止搜索

generate_final_response：把最优路径的完整规划交给 LLM，生成最终答案

_run_node 原理：

queue = deque([A])
第一轮：
popleft → A，队列空
A生成子B,C,D
反转 [B,C,D] → [D,C,B]
依次appendleft → queue = [B,C,D]

第二轮：
popleft → B，队列变成 [C,D]
B生成子B1,B2
反转 [B1,B2] → [B2,B1]
appendleft → queue = [B1,B2,C,D]

第三轮：
popleft → B1，队列 [B2,C,D]
......

一直往最深的分支钻，直到某个节点没有子步骤(next_steps 是空列表)

ToT 和普通 ReAct 的区别