LangGraph Tree of Thought：让 AI 多想几个方案再决定，还得算清这笔 Token 账

目录

一、为什么需要 Tree of Thought：重要决策，本来就不该“一条路走到黑”

二、先把误区说清：ToT 不是“一个问题问模型三遍”

（一）线性推理与本文 ToT 的差异

（二）全文核心代码展示

三、核心代码展开详细讲解

（一）整条链路：8 个节点、3 个条件路由

每个节点分别做什么？

（二）状态设计：path / frontier / pruned 才是 ToT 真正的“记忆”

1. 搜索上下文：决定“下一步从哪里继续想”

2. 预算闸门：决定“最多能烧多少”

3. 可观测性：决定“能不能复盘与算账”

（三）主要的处理逻辑详细说明

1. 第一把省钱刀：批量评估，把 N 次调用压成 1 次

2. 第二把省钱刀：剪枝——砍掉一个节点，等于砍掉整棵子树

一个很容易写错的地方：frontier 应该收谁？

3. 第三把省钱刀：早停——“够好”不等于“第一层高分就停”

4. 第四把省钱刀：记忆化缓存——回溯与重跑不再重复付费

（四）实战代码：图结构本身反而是最简单的部分

四、运行实记与演示说明

（一）运行方式

（二）具体演示结果详细说明

1. 运行演示 1：三层决策真的串起来了

第一步：先由 LLM 拆成 3 个逐层目标

第二步：第一层批量评估

第三步：第二层候选开始“带着上一层答案思考”

这次搜索一共花了多少？

2. 运行演示 2：撞上死胡同，ToT 才真正体现价值

回溯不是免费的

3. 运行演示 3：深度 + 展开预算，给搜索空间上“双保险”

4. 运行演示 4：早停 + 缓存，两种省法可以叠加

（三）把账算到图上：优化后的调用次数到底省了多少？

五、问题排查与线上注意事项

（一）六个最容易踩的坑，以及应该怎么排查

坑 1：早停阈值太低，多级链路第一层就结束

坑 2：评估维度名对不上，维度分悄悄丢失

坑 3：备选池一直为空，回溯逻辑实际上是死代码

坑 4：缓存键漏掉 path，回溯之后生成一模一样的候选

坑 5：一调大深度和分支数，账单和运行时间一起爆炸

坑 6：每个候选单独打分，结果全部 8 分上下

（二）生产级还应该补什么？

1. 并行展开

2. 跨会话共享缓存

3. frontier 持久化

4. 发散模型和评估模型分开

5. 和其他推理模式组合

（三）什么时候值得用 ToT，什么时候不值得？

适合使用

不太适合

六、Stage 4 收官：真正能上线的推理 Agent，都必须有“硬边界”

附：本文核心参数速查

干货分享，感谢您的阅读！

这是「LangGraph Agent Engineering Mastery」系列  Stage 4 推理 Agent · 第 5 篇（Stage 4 收官）。 这不是一篇“让模型多生成几个答案”的概念介绍，而是一篇真正从工程链路、搜索空间、回溯状态和调用成本出发的 Tree of Thought（ToT）实战。

读完你会得到什么？

• 用 LangGraph 搭出一条真正的多级 ToT 链路：分解 → 发散 → 批量评估 → 束搜索剪枝 → 选优下探 → 回溯 → 整合。

• 明白 decompose 为什么是“多级决策”与“同一个问题问三遍”的分水岭。

• 看懂 path / frontier / pruned 三组状态怎样共同支撑搜索与回溯。

• 掌握四个最实用的成本控制手段：批量评估、剪枝、早停、记忆化缓存。

• 通过 4 组真实运行演示，直观看到：三层上下文如何传递、死胡同如何回溯、预算怎样止损、缓存怎样把重复调用降到 0。

ToT 真正难的从来不是“让 AI 多想几个方案”，而是让它在有限 Token、有限时间和有限搜索深度里，系统地比较、舍弃、换路，并且能够解释自己为什么停下来。

一、为什么需要 Tree of Thought：重要决策，本来就不该“一条路走到黑”

做一个重要技术选型，比如：“高并发订单系统到底应该用什么架构？”你大概率不会想到一个方案就立刻拍板。更符合真实工程决策的过程通常是：

这里最关键的一点是：后一层的问题，是前一层决策成立之后才存在的。

你不会在“架构模式”都没决定时，就提前争论“分库分表究竟分 64 张还是 128 张”。这也是 Tree of Thought 和普通“多采样”最本质的区别。

前面的 ReAct、Planning、Reflection、Self-Correction，虽然推理方式不同，但大体仍沿着一条主路径前进。ToT 则把搜索算法带进了推理过程：每层发散多个候选、统一评估、主动剪枝、选择最优继续深入，必要时还能回溯。

如果你做过后端，这个过程很像数据库查询优化器：

查询优化器不会把所有可能计划完整执行一遍，因为那样成本不可接受。ToT 也是一样：如果没有剪枝和预算，分支数 × 深度会很快把 Token 成本推向指数级。

二、先把误区说清：ToT 不是“一个问题问模型三遍”

很多 ToT Demo 看起来像这样：

同一个问题生成 3 个答案 → 分别打分 → 选最高分。

这当然算“多候选推理”，但还不是本文强调的多级 Tree of Thought。

（一）线性推理与本文 ToT 的差异

真正的多级 ToT，至少要满足三件事：

所以本文先用 decompose 节点，把一个大任务拆成逐层决策目标，再由 expand 带着“当前层目标 + 已选路径 + 已剪枝方向”继续发散。

这一步非常关键。没有它，你只是做了“多次回答”；有了它，才真正形成一棵逐层依赖、可以搜索和回溯的决策树。

线性推理与本文 ToT 的差异：

（二）全文核心代码展示

基于这个大前提，我们准备本次教学的基本代码整体如下：

"""Demo 05: Tree of Thought — 多级链路的分层发散、批量评估、束搜索剪枝与回溯。

演示 Tree of Thought 推理模式（全链路真实 LLM 调用）：
1. Decompose 节点：真实 LLM 把任务拆成"逐层决策目标"和统一评估维度
2. Expand 节点：真实 LLM 在当前层生成互不同质的候选思路（承接已选路径）
3. Evaluate 节点：真实 LLM 一次调用批量给整层候选打分（多维度 + 可行性）
4. Prune 节点：束搜索剪枝，保留 top-K，落选可行候选进入 frontier 备用
5. Select 节点：选出本层最优推进深度，质量达标时提前收敛
6. Backtrack 节点：本层全军覆没时，从 frontier 回到全局次优候选换条路走
7. Synthesize 节点：真实 LLM 把整条决策路径整合成可落地的最终方案

链路效率优化（Token 是朴素 ToT 的痛点，这里给出四种可组合的解法）：
– 批量评估：整层候选合并为 1 次评估调用，把 N 次调用压成 1 次
– 束搜索 + 阈值剪枝：低分/不可行候选不再向下展开，砍掉整棵子树的开销
– 提前收敛：本层最优分达到 satisfaction_score 即停止下探
– 记忆化缓存：相同上下文的展开/评估复用结果，回溯时不重复烧 Token

本 Demo 通过 shared.get_llm() 调用 .env 配置的真实在线模型
（fallback_to_mock=False，需联网）。LLM 输出无法解析时回退到启发式保底
逻辑（也用于 mock/离线环境，保证流程可运行）。

运行方式：
python stages/stage4_reasoning/05_tree_of_thought/main.py
"""

from __future__ import annotations

import hashlib
import json
import operator
import re
import sys
import time
from functools import partial
from pathlib import Path
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent))

from shared import get_llm, get_logger, log_step, log_success, log_warning

logger = get_logger("demo.04_05_tree_of_thought")

MAX_BRANCHES = 3
MAX_DEPTH = 3
BEAM_WIDTH = 2
SCORE_THRESHOLD = 5.0
SATISFACTION_SCORE = 9.5
MIN_EXPLORED_DEPTH = 2
MAX_BACKTRACKS = 2
EXPANSION_BUDGET = 6
FRONTIER_CAPACITY = 6

DEFAULT_CRITERIA = ["可行性", "性能", "成本", "风险"]

# 死胡同注入开关（仅供"剪枝与回溯"演示使用）：指定层的候选全部判为不可行，
# 模拟真实场景中某层方案集体撞上硬约束。一次性生效，触发后自动清除
_DEAD_END_AT_DEPTH: int | None = None

# ============================================================
# 真实 LLM（通过 shared.get_llm 获取 .env 配置的在线模型）
# ============================================================
_LLM = None

# 记忆化缓存：key = 节点类型 + 上下文签名。回溯重新展开同一层、
# 或重复运行同一任务时直接复用，不产生真实调用
_LLM_CACHE: dict[str, str] = {}
_CALL_STATS: dict[str, int] = {"llm_calls": 0, "cache_hits": 0}

def _get_tot_llm():
"""获取真实在线 LLM 实例（模块内复用，fallback_to_mock=False 确保真实调用）。"""
global _LLM
if _LLM is None:
_LLM = get_llm(fallback_to_mock=False)
return _LLM

def _cache_key(kind: str, *parts: Any) -> str:
"""按调用类型 + 上下文生成缓存键。"""
raw = json.dumps([kind, *parts], ensure_ascii=False, sort_keys=True, default=str)
return f"{kind}:{hashlib.sha1(raw.encode('utf-8')).hexdigest()}"

def _cached_invoke(key: str, messages: list[BaseMessage]) -> str:
"""带记忆化的 LLM 调用：命中缓存时不发起真实请求。"""
if key in _LLM_CACHE:
_CALL_STATS["cache_hits"] += 1
return _LLM_CACHE[key]

response = _get_tot_llm().invoke(messages)
text = str(response.content).strip()
_LLM_CACHE[key] = text
_CALL_STATS["llm_calls"] += 1
return text

def _reset_stats() -> None:
"""重置调用统计（每个演示单独统计，缓存本身跨演示保留）。"""
_CALL_STATS["llm_calls"] = 0
_CALL_STATS["cache_hits"] = 0

def _parse_llm_json(text: str) -> dict | None:
"""从 LLM 输出中解析 JSON（容忍 Markdown 代码块包裹等格式噪音）。"""
text = text.strip()
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\\s*|\\s*```$", "", text, flags=re.S).strip()
try:
parsed = json.loads(text)
return parsed if isinstance(parsed, dict) else None
except json.JSONDecodeError:
match = re.search(r"\\{.*\\}", text, re.S)
if match:
try:
parsed = json.loads(match.group())
return parsed if isinstance(parsed, dict) else None
except json.JSONDecodeError:
return None
return None

def _trace_event(node: str, detail: str) -> dict:
"""构造一条执行轨迹事件。"""
return {"node": node, "detail": detail, "ts": time.strftime("%H:%M:%S")}

# ============================================================
# Tree of Thought State
# ============================================================
class ThoughtBranch(TypedDict, total=False):
id: str
description: str
reasoning: str
score: float
is_viable: bool
depth: int
parent_id: str | None
dimensions: dict[str, float]
risk: str
comment: str

class TreeOfThoughtState(TypedDict, total=False):
messages: Annotated[list[BaseMessage], add_messages]
task: str
branches: list[ThoughtBranch]
selected_branch: ThoughtBranch | None
exploration_history: list[dict]
current_depth: int
max_depth: int
max_branches: int
best_solution: str
backtrack_count: int
# 分层搜索上下文
stage_goals: list[str] # 每一层要回答的决策问题（LLM 拆解）
criteria: list[str] # 统一评估维度（LLM 拆解）
path: list[ThoughtBranch] # 已选定的决策路径（根 → 当前）
beam: list[ThoughtBranch] # 本层剪枝后存活的候选
frontier: list[ThoughtBranch] # 全局备选池（回溯用）
pruned: list[dict] # 被剪掉的候选（供后续层避免重复思路）
# 搜索预算与阈值
beam_width: int
score_threshold: float
satisfaction_score: float
max_backtracks: int
expansion_budget: int
# 结果与可观测性
final_solution: str
stop_reason: str
metrics: dict
execution_trace: Annotated[list[dict], operator.add]

_DEFAULT_METRICS: dict[str, int] = {
"llm_calls": 0,
"cache_hits": 0,
"expansions": 0,
"generated": 0,
"pruned": 0,
"naive_calls": 0,
}

def _conf(state: TreeOfThoughtState, defaults: dict | None, key: str, fallback: Any) -> Any:
"""读取搜索参数：state 优先，其次建图时绑定的默认值，最后模块常量。"""
value = state.get(key)
if value is None:
value = (defaults or {}).get(key, fallback)
return value

def _bump_metrics(state: TreeOfThoughtState, before: dict[str, int], **delta: int) -> dict:
"""累计效率指标：LLM 调用/缓存命中取自全局统计的增量。"""
metrics = {**_DEFAULT_METRICS, **state.get("metrics", {})}
metrics["llm_calls"] += _CALL_STATS["llm_calls"] – before["llm_calls"]
metrics["cache_hits"] += _CALL_STATS["cache_hits"] – before["cache_hits"]
for key, value in delta.items():
metrics[key] = metrics.get(key, 0) + value
return metrics

def _path_signature(path: list[ThoughtBranch]) -> list[str]:
"""决策路径的紧凑签名（用作缓存键的一部分）。"""
return [str(node.get("description", "")) for node in path]

def _prune_reason(branch: ThoughtBranch, threshold: float, dead_end: bool) -> str:
"""给出候选被剪枝的原因（写进 pruned，后续层据此避免重复发散）。"""
if dead_end:
return "违反硬约束"
if not branch.get("is_viable", True):
return "不可行"
return f"分数低于阈值 {threshold}"

def _format_path(path: list[ThoughtBranch]) -> str:
"""把已选路径渲染为提示词片段。"""
if not path:
return "（这是第一层，尚未做出任何决策）"
return "\\n".join(
f"第 {i + 1} 层已选: {node.get('description', '')}（{node.get('reasoning', '')[:60]}）"
for i, node in enumerate(path)
)

# ============================================================
# 启发式保底逻辑（LLM 输出无法解析时使用，也用于 mock/离线环境）
# ============================================================
def _fallback_branches(task: str, depth: int, max_branches: int) -> list[ThoughtBranch]:
"""启发式模板分支（保底方案）。"""
if "设计" in task or "架构" in task:
branches: list[ThoughtBranch] = [
{
"id": f"d{depth}_a",
"description": "微服务架构方案",
"reasoning": "将系统拆分为独立微服务，每个服务独立部署。优点：扩展性好、故障隔离。缺点：复杂度高。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_b",
"description": "单体架构 + 模块化方案",
"reasoning": "保持单体但严格模块化，通过包结构隔离。优点：简单、部署方便。缺点：扩展性有限。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_c",
"description": "Serverless 方案",
"reasoning": "使用云函数实现核心逻辑。优点：免运维、按需付费。缺点：冷启动、调试困难。",
"score": 0.0,
"is_viable": True,
},
]
elif "解决" in task or "问题" in task or "优化" in task:
branches = [
{
"id": f"d{depth}_a",
"description": "渐进式优化方案",
"reasoning": "分阶段逐步优化，每阶段验证效果。风险低，见效慢。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_b",
"description": "根因分析+一次性修复",
"reasoning": "深入分析根本原因，一次性彻底修复。风险高，但见效快。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_c",
"description": "替换方案",
"reasoning": "用全新方案替换现有实现。工作量大，但可避免历史包袱。",
"score": 0.0,
"is_viable": depth == 0,
},
]
else:
branches = [
{
"id": f"d{depth}_a",
"description": "方案 A：保守稳健路线",
"reasoning": "采用成熟技术，优先保证稳定性。适合生产环境。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_b",
"description": "方案 B：创新激进路线",
"reasoning": "采用前沿技术，追求最优性能。适合技术探索。",
"score": 0.0,
"is_viable": True,
},
{
"id": f"d{depth}_c",
"description": "方案 C：折中平衡路线",
"reasoning": "核心模块用成熟技术，非核心模块尝试新技术。兼顾稳定和创新。",
"score": 0.0,
"is_viable": True,
},
]

return branches[:max_branches]

def score_branch(branch: ThoughtBranch, task: str, depth: int) -> float:
"""启发式评分（保底方案）：按分支序位给基准分，深度加成，不可行重罚。"""
base_scores = {
"a": 7.5,
"b": 6.5,
"c": 5.0,
}

branch_letter = str(branch.get("id", ""))[-1:]
base = base_scores.get(branch_letter, 5.0)
depth_bonus = depth * 0.5
viability_penalty = 0 if branch.get("is_viable", True) else -3.0

return max(0.0, min(round(base + depth_bonus + viability_penalty, 1), 10.0))

# ============================================================
# LLM 分解：把任务拆成逐层决策目标 + 统一评估维度
# ============================================================
_DECOMPOSE_SYSTEM_PROMPT = """你是复杂问题分解专家，负责为「思维树（Tree of Thought）」搜索规划分层决策。

只输出一个 JSON 对象（不要输出任何其他文字）：
{"stage_goals": ["第 1 层要决定的问题", "第 2 层要决定的问题", …],
"criteria": ["评估维度1", "评估维度2", "评估维度3"]}

要求：
1. stage_goals 按"从宏观到细节"排列，每层只解决一个关键决策点，数量等于给定层数
2. 后一层必须是在前一层决策已确定的前提下才需要回答的问题
3. criteria 为 3~4 个给候选方案打分的维度（如 可行性 / 性能 / 成本 / 风险）"""

def decompose_task(task: str, max_depth: int) -> tuple[list[str], list[str]]:
"""调用真实 LLM 把任务拆成逐层决策目标与评估维度；解析失败时回退到通用模板。"""
key = _cache_key("decompose", task, max_depth)
text = _cached_invoke(
key,
[
SystemMessage(content=_DECOMPOSE_SYSTEM_PROMPT),
HumanMessage(content=f"任务: {task}\\n搜索层数: {max_depth}\\n请输出 JSON。"),
],
)
decision = _parse_llm_json(text)

goals: list[str] = []
criteria: list[str] = []
if decision:
goals = [str(g).strip() for g in decision.get("stage_goals", []) if str(g).strip()]
criteria = [str(c).strip() for c in decision.get("criteria", []) if str(c).strip()]

if not goals:
log_warning(logger, "LLM 分层目标输出无法解析为 JSON，使用保底通用目标")
goals = [f"第 {i + 1} 层：确定该层的关键决策" for i in range(max_depth)]
if not criteria:
criteria = list(DEFAULT_CRITERIA)

goals = goals[:max_depth]
while len(goals) < max_depth:
goals.append(f"第 {len(goals) + 1} 层：在已有决策上继续细化")

return goals, criteria[:4]

# ============================================================
# LLM 发散：在当前层生成候选思路
# ============================================================
_EXPAND_SYSTEM_PROMPT = """你是方案发散专家，负责在思维树的某一层生成互不同质的候选思路。

只输出一个 JSON 对象（不要输出任何其他文字）：
{"branches": [{"description": "<一句话方案名>",
"reasoning": "<关键思路 + 优点 + 缺点>",
"risk": "<主要风险>",
"is_viable": true}]}

要求：
1. 候选之间思路差异明显，禁止同质化换皮
2. 必须承接「已选路径」，回答的是"本层决策目标"，不要重新回答整个任务
3. 不要重复「已剪枝思路」中列出的方向
4. 候选数量不超过给定上限"""

def generate_thought_branches(
task: str,
depth: int,
history: list,
goal: str = "",
path: list[ThoughtBranch] | None = None,
pruned: list[dict] | None = None,
max_branches: int = MAX_BRANCHES,
) -> list[ThoughtBranch]:
"""调用真实 LLM 生成当前层的候选思路；解析失败时回退到启发式模板分支。

Args:
task: 原始任务
depth: 当前层（从 0 开始）
history: 已确定的决策历程（用于提示词上下文）
goal: 本层决策目标（由 decompose 阶段规划）
path: 已选定的决策路径
pruned: 之前被剪掉的思路，避免重复发散
max_branches: 本层候选数上限
"""
path = list(path or [])
pruned = list(pruned or [])

pruned_text = (
"\\n".join(f"- {p.get('description', '')}（{p.get('reason', '')}）" for p in pruned[-5:])
if pruned
else "（暂无）"
)
key = _cache_key(
"expand", task, depth, goal, _path_signature(path), max_branches, len(pruned)
)
text = _cached_invoke(
key,
[
SystemMessage(content=_EXPAND_SYSTEM_PROMPT),
HumanMessage(
content=(
f"原始任务: {task}\\n"
f"当前层: 第 {depth + 1} 层\\n"
f"本层决策目标: {goal or '在已有决策上给出下一步方案'}\\n"
f"已选路径:\\n{_format_path(path)}\\n"
f"已剪枝思路:\\n{pruned_text}\\n"
f"候选数上限: {max_branches}\\n"
f"请输出候选 JSON。"
)
),
],
)
decision = _parse_llm_json(text)

branches: list[ThoughtBranch] = []
raw_items = decision.get("branches", []) if decision else []
if isinstance(raw_items, list):
for i, item in enumerate(raw_items[:max_branches]):
if not isinstance(item, dict):
continue
description = str(item.get("description", "")).strip()
if not description:
continue
branches.append({
"id": f"d{depth}_{chr(ord('a') + i)}",
"description": description,
"reasoning": str(item.get("reasoning", "")).strip(),
"risk": str(item.get("risk", "")).strip(),
"score": 0.0,
"is_viable": bool(item.get("is_viable", True)),
})

if not branches:
log_warning(logger, "LLM 候选输出无法解析为 JSON，使用启发式保底分支")
branches = _fallback_branches(task, depth, max_branches)

parent_id = path[-1].get("id") if path else None
for branch in branches:
branch["depth"] = depth
branch["parent_id"] = parent_id

return branches[:max_branches]

# ============================================================
# LLM 批量评估：一次调用给整层候选打分
# ============================================================
_EVALUATE_SYSTEM_PROMPT = """你是严格的方案评估专家，一次性评估同一层的全部候选方案。

只输出一个 JSON 对象（不要输出任何其他文字）：
{"scores": [{"id": "<候选ID>", "score": <0-10 数字>,
"dimensions": {"<维度名>": <0-10 数字>},
"is_viable": true, "comment": "<一句话理由>"}]}

要求：
1. 必须为每个候选给出一条记录，id 与输入完全一致
2. score 为综合分，评分严格并拉开差距，禁止所有候选同分
3. dimensions 使用给定的评估维度
4. 违反约束、无法落地的候选 is_viable 置为 false"""

def score_branches(
task: str,
branches: list[ThoughtBranch],
depth: int,
criteria: list[str],
goal: str = "",
) -> tuple[list[ThoughtBranch], str]:
"""调用真实 LLM 批量评估整层候选，返回 (已打分候选, 来源标记 llm/fallback)。

整层合并为一次调用是本 Demo 的核心效率优化：朴素 ToT 每个候选一次评估调用，
这里 N 个候选只花 1 次调用，且模型能横向比较、更容易拉开分差。
"""
if not branches:
return [], "fallback"

candidates_text = "\\n".join(
f"- id={b['id']} | 方案: {b['description']} | 思路: {b.get('reasoning', '')[:120]}"
f" | 风险: {b.get('risk', '') or '未说明'}"
for b in branches
)
key = _cache_key(
"evaluate", task, depth, goal, criteria, [(b["id"], b["description"]) for b in branches]
)
text = _cached_invoke(
key,
[
SystemMessage(content=_EVALUATE_SYSTEM_PROMPT),
HumanMessage(
content=(
f"原始任务: {task}\\n"
f"本层决策目标: {goal or '选出最优的下一步决策'}\\n"
f"评估维度: {', '.join(criteria)}\\n"
f"候选方案:\\n{candidates_text}\\n\\n"
f"请输出评分 JSON。"
)
),
],
)
decision = _parse_llm_json(text)
score_map = _index_scores(decision, criteria)

scored: list[ThoughtBranch] = []
source = "llm" if score_map else "fallback"
if not score_map:
log_warning(logger, "LLM 评分输出无法解析为 JSON，使用启发式保底评分")

for branch in branches:
item = dict(branch)
entry = score_map.get(branch["id"])
if entry:
item["score"] = entry["score"]
item["dimensions"] = entry["dimensions"]
item["comment"] = entry["comment"]
item["is_viable"] = entry["is_viable"] and branch.get("is_viable", True)
else:
item["score"] = score_branch(branch, task, depth)
item["dimensions"] = {}
item["comment"] = item.get("comment", "")
scored.append(item)

scored.sort(key=lambda b: b["score"], reverse=True)
return scored, source

def _index_scores(decision: dict | None, criteria: list[str]) -> dict[str, dict]:
"""把 LLM 评分 JSON 规整为 {候选ID: 评分记录}；非法记录直接丢弃走保底。"""
if not decision or not isinstance(decision.get("scores"), list):
return {}

indexed: dict[str, dict] = {}
for item in decision["scores"]:
if not isinstance(item, dict):
continue
branch_id = str(item.get("id", "")).strip()
if not branch_id:
continue
try:
score = float(item["score"])
except (KeyError, TypeError, ValueError):
continue

raw_dims = item.get("dimensions", {})
dimensions: dict[str, float] = {}
if isinstance(raw_dims, dict):
for name in criteria:
try:
dimensions[name] = max(0.0, min(float(raw_dims[name]), 10.0))
except (KeyError, TypeError, ValueError):
continue

indexed[branch_id] = {
"score": max(0.0, min(round(score, 1), 10.0)),
"dimensions": dimensions,
"is_viable": bool(item.get("is_viable", True)),
"comment": str(item.get("comment", "")).strip(),
}
return indexed

# ============================================================
# LLM 整合：把决策路径写成最终方案
# ============================================================
_SYNTHESIZE_SYSTEM_PROMPT = """你是方案落地专家，负责把思维树搜索选出的决策路径整合成最终方案。

输出 Markdown 正文，包含四部分：
1. 最终方案概述
2. 逐层决策理由（说明每层为什么选它、放弃了什么）
3. 落地要点
4. 主要风险与应对

不要输出 JSON，不要复述提示词。"""

def synthesize_solution(task: str, path: list[ThoughtBranch], criteria: list[str]) -> str:
"""调用真实 LLM 把决策路径整合为最终方案；输出不可用时回退到路径拼接。"""
fallback = "\\n".join(
f"{i + 1}. {node.get('description', '')} — {node.get('reasoning', '')}"
for i, node in enumerate(path)
)
if not path:
return fallback

key = _cache_key("synthesize", task, _path_signature(path))
text = _cached_invoke(
key,
[
SystemMessage(content=_SYNTHESIZE_SYSTEM_PROMPT),
HumanMessage(
content=(
f"原始任务: {task}\\n"
f"评估维度: {', '.join(criteria)}\\n"
f"最终决策路径:\\n"
+ "\\n".join(
f"第 {i + 1} 层: {node.get('description', '')}"
f"（得分 {node.get('score', 0)}）— {node.get('reasoning', '')}"
for i, node in enumerate(path)
)
+ "\\n\\n请输出最终方案。"
)
),
],
)
if len(text) < 30:
log_warning(logger, "LLM 整合输出过短，使用决策路径拼接保底")
return fallback
return text

# ============================================================
# 节点实现
# ============================================================
def decompose_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Decompose 节点：真实 LLM 规划逐层决策目标与统一评估维度。"""
before = dict(_CALL_STATS)
task = state["task"]
max_depth = _conf(state, defaults, "max_depth", MAX_DEPTH)

log_step(logger, "Decompose", f"规划 {max_depth} 层决策目标（真实 LLM）")
print(f"\\n [分解] 任务: {task}")

goals, criteria = decompose_task(task, max_depth)

for i, goal in enumerate(goals):
print(f" 第 {i + 1} 层目标: {goal}")
print(f" 评估维度: {' / '.join(criteria)}")

return {
"stage_goals": goals,
"criteria": criteria,
"path": [],
"frontier": [],
"pruned": [],
"beam": [],
"metrics": _bump_metrics(state, before, naive_calls=1),
"execution_trace": [_trace_event("decompose", f"{len(goals)} 层目标 / {len(criteria)} 个维度")],
}

def expand_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Expand 节点：真实 LLM 在当前层生成候选思路（承接已选路径，避开已剪枝方向）。"""
before = dict(_CALL_STATS)
task = state["task"]
depth = state.get("current_depth", 0)
max_branches = _conf(state, defaults, "max_branches", MAX_BRANCHES)
goals = state.get("stage_goals", [])
goal = goals[depth] if depth < len(goals) else ""

log_step(logger, "Expand", f"第 {depth + 1} 层发散候选（真实 LLM）")
print(f"\\n [发散] 第 {depth + 1} 层「{goal or '下一步决策'}」生成候选:")

branches = generate_thought_branches(
task,
depth,
state.get("exploration_history", []),
goal=goal,
path=state.get("path", []),
pruned=state.get("pruned", []),
max_branches=max_branches,
)

for b in branches:
print(f" 候选 [{b['id']}] {b['description']}")
print(f" 思路: {b.get('reasoning', '')[:70]}…")

budget = _conf(state, defaults, "expansion_budget", EXPANSION_BUDGET)
metrics = _bump_metrics(
state, before, expansions=1, generated=len(branches), naive_calls=1 + len(branches)
)

return {
"branches": branches,
"expansion_budget": budget – 1,
"metrics": metrics,
"execution_trace": [_trace_event("expand", f"第 {depth + 1} 层生成 {len(branches)} 个候选")],
}

def evaluate_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Evaluate 节点：真实 LLM 一次调用批量给整层候选打分。"""
before = dict(_CALL_STATS)
task = state["task"]
depth = state.get("current_depth", 0)
branches = state.get("branches", [])
criteria = state.get("criteria") or list(DEFAULT_CRITERIA)
goals = state.get("stage_goals", [])
goal = goals[depth] if depth < len(goals) else ""

log_step(logger, "Evaluate", f"批量评估 {len(branches)} 个候选（1 次调用）")
print(f"\\n [评估] 一次调用批量打分（{len(branches)} 个候选）:")

scored, source = score_branches(task, branches, depth, criteria, goal)

for b in scored:
status = "✓" if b.get("is_viable", True) else "✗"
dims = b.get("dimensions") or {}
dims_text = " ".join(f"{k} {v}" for k, v in dims.items())
print(f" [{status}] {b['id']} {b['description']} → {b['score']} 分 {dims_text}")
if b.get("comment"):
print(f" 理由: {b['comment'][:70]}")

return {
"branches": scored,
"metrics": _bump_metrics(state, before),
"execution_trace": [
_trace_event("evaluate", f"第 {depth + 1} 层批量评分（来源: {source}）")
],
}

def prune_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Prune 节点：束搜索剪枝，每层只保留 beam_width 条候选存活。

存活的候选里第 1 名立即下探，其余进 frontier 等回溯；束宽之外以及不可行、
低于阈值的候选一律淘汰。剪枝是 ToT 最有效的省 Token 手段：砍掉一个候选，
等于砍掉它下面整棵子树的展开与评估开销。
"""
global _DEAD_END_AT_DEPTH

before = dict(_CALL_STATS)
branches = state.get("branches", [])
depth = state.get("current_depth", 0)
beam_width = _conf(state, defaults, "beam_width", BEAM_WIDTH)
threshold = _conf(state, defaults, "score_threshold", SCORE_THRESHOLD)

dead_end = _DEAD_END_AT_DEPTH is not None and depth == _DEAD_END_AT_DEPTH
if dead_end:
_DEAD_END_AT_DEPTH = None
log_warning(logger, f"第 {depth + 1} 层候选集体撞上硬约束（模拟），本层全部剪枝")
print(f"\\n [剪枝] ⚠ 第 {depth + 1} 层候选全部违反硬约束（模拟死胡同）")

survivors = (
[]
if dead_end
else [b for b in branches if b.get("is_viable", True) and b.get("score", 0) >= threshold]
)
survivors.sort(key=lambda b: b.get("score", 0), reverse=True)

# beam[0] 立即下探，beam[1:] 进备选池等回溯；束宽之外的候选直接淘汰
beam = survivors[:beam_width]
alternates = beam[1:]
dropped = [
(b, _prune_reason(b, threshold, dead_end)) for b in branches if b not in survivors
]
dropped += [(b, f"超出束宽 {beam_width}") for b in survivors[beam_width:]]

log_step(logger, "Prune", f"保留 {len(beam)}/{len(branches)}，剪枝 {len(dropped)}")
print(f"\\n [剪枝] 束宽 {beam_width}，阈值 {threshold}：")
for b in beam:
print(f" ✓ 保留 [{b['id']}] {b['description']}（{b.get('score', 0)} 分）")
for b, reason in dropped:
print(f" ✗ 剪枝 [{b['id']}] {b['description']}（{reason}）")

pruned = list(state.get("pruned", []))
pruned += [
{
"id": b["id"],
"depth": depth,
"description": b["description"],
"reason": reason,
}
for b, reason in dropped
]

frontier = list(state.get("frontier", [])) + alternates
frontier.sort(key=lambda b: b.get("score", 0), reverse=True)
frontier = frontier[:FRONTIER_CAPACITY]

if alternates:
print(f" ○ 备选入池 {len(alternates)} 个（回溯时可用），当前池容量 {len(frontier)}")

return {
"beam": beam,
"branches": beam,
"pruned": pruned,
"frontier": frontier,
"metrics": _bump_metrics(state, before, pruned=len(dropped)),
"execution_trace": [
_trace_event("prune", f"第 {depth + 1} 层保留 {len(beam)}，剪枝 {len(dropped)}")
],
}

def select_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Select 节点：选出本层最优推进深度；质量达标时标记提前收敛。"""
beam = state.get("beam", [])
depth = state.get("current_depth", 0)
satisfaction = _conf(state, defaults, "satisfaction_score", SATISFACTION_SCORE)
max_depth = _conf(state, defaults, "max_depth", MAX_DEPTH)
history = list(state.get("exploration_history", []))
path = list(state.get("path", []))

if not beam:
log_warning(logger, "本层无存活候选，交由回溯处理")
return {"selected_branch": None}

best = beam[0]
log_step(logger, "Select", f"选定 {best['id']}（{best.get('score', 0)} 分）")
print(f"\\n [选择] ★ 第 {depth + 1} 层选定: {best['description']}（{best.get('score', 0)} 分）")

path.append(best)
history.append({
"depth": depth,
"selected": best["id"],
"description": best["description"],
"score": best.get("score", 0),
"alternatives": [b["id"] for b in beam[1:]],
"via_backtrack": False,
})

# 早停要求"分数足够高"且"已探索足够层数"：只凭第一层的高分就收敛，
# 等于放弃了分层决策本身的价值
stop_reason = ""
min_depth = min(MIN_EXPLORED_DEPTH, max_depth)
if best.get("score", 0) >= satisfaction and depth + 1 >= min_depth:
stop_reason = f"提前收敛（{best.get('score', 0)} 分 ≥ 满意阈值 {satisfaction}）"
print(f" [早停] {stop_reason}，跳过后续层的发散与评估")

return {
"selected_branch": best,
"path": path,
"exploration_history": history,
"current_depth": depth + 1,
"best_solution": f"{best['description']} — {best.get('reasoning', '')}",
"stop_reason": stop_reason,
"execution_trace": [
_trace_event("select", f"第 {depth + 1} 层选定 {best['id']}（{best.get('score', 0)} 分）")
],
}

def backtrack_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Backtrack 节点：本层全军覆没时，回到 frontier 中的全局次优候选换条路走。"""
frontier = list(state.get("frontier", []))
history = list(state.get("exploration_history", []))
backtrack_count = state.get("backtrack_count", 0)

candidate = frontier.pop(0)
fallback_depth = candidate.get("depth", 0)
path = list(state.get("path", []))[:fallback_depth]
path.append(candidate)

log_warning(
logger,
f"回溯到第 {fallback_depth + 1} 层备选 {candidate['id']}（{candidate.get('score', 0)} 分）",
)
print(
f"\\n [回溯] ↩ 本层无可行候选，退回第 {fallback_depth + 1} 层备选: "
f"{candidate['description']}（{candidate.get('score', 0)} 分）"
)

history.append({
"depth": fallback_depth,
"selected": candidate["id"],
"description": candidate["description"],
"score": candidate.get("score", 0),
"alternatives": [b["id"] for b in frontier],
"via_backtrack": True,
})

return {
"selected_branch": candidate,
"path": path,
"frontier": frontier,
"exploration_history": history,
"current_depth": fallback_depth + 1,
"backtrack_count": backtrack_count + 1,
"best_solution": f"{candidate['description']} — {candidate.get('reasoning', '')}",
"execution_trace": [
_trace_event("backtrack", f"回溯到 {candidate['id']}（第 {fallback_depth + 1} 层）")
],
}

def synthesize_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Synthesize 节点：真实 LLM 把整条决策路径整合成最终方案。"""
before = dict(_CALL_STATS)
task = state["task"]
path = state.get("path", [])
criteria = state.get("criteria") or list(DEFAULT_CRITERIA)

log_step(logger, "Synthesize", f"整合 {len(path)} 层决策为最终方案（真实 LLM）")
print(f"\\n [整合] 把 {len(path)} 层决策整合为最终方案…")

solution = synthesize_solution(task, path, criteria)
print(f" [方案] {solution[:100]}{'…' if len(solution) > 100 else ''}")

return {
"final_solution": solution,
"metrics": _bump_metrics(state, before, naive_calls=1),
"execution_trace": [_trace_event("synthesize", f"生成最终方案（{len(solution)} 字）")],
}

def output_node(state: TreeOfThoughtState, defaults: dict | None = None) -> dict:
"""Output 节点：输出决策路径、效率指标与最终方案。"""
history = state.get("exploration_history", [])
metrics = {**_DEFAULT_METRICS, **state.get("metrics", {})}
backtrack_count = state.get("backtrack_count", 0)
stop_reason = state.get("stop_reason") or "达到搜索边界"
solution = state.get("final_solution") or state.get("best_solution", "")

log_success(
logger,
f"搜索完成，深度 {state.get('current_depth', 0)}，回溯 {backtrack_count} 次，"
f"真实调用 {metrics['llm_calls']} 次",
)
print("\\n [完成] 思维树搜索结束!")

summary_parts = [
f"最终方案:\\n{solution}\\n",
"— 探索历程 —",
f"搜索深度: {state.get('current_depth', 0)} | 回溯次数: {backtrack_count} | 结束原因: {stop_reason}",
"决策路径:",
]
for h in history:
alts = ", ".join(h.get("alternatives", [])) or "无"
marker = "↩" if h.get("via_backtrack") else "→"
summary_parts.append(
f" {marker} 第 {h['depth'] + 1} 层: [{h['selected']}] {h['description']} "
f"(分数: {h['score']}, 备选: {alts})"
)

summary_parts += [
"",
"— 效率指标 —",
f"真实 LLM 调用: {metrics['llm_calls']} 次 | 缓存命中: {metrics['cache_hits']} 次",
f"展开层数: {metrics['expansions']} | 生成候选: {metrics['generated']} | 剪枝: {metrics['pruned']}",
f"朴素 ToT 预计调用: {metrics['naive_calls']} 次（每候选单独评估）",
]

print(f" 结束原因: {stop_reason}")
print(
f" 效率: 真实调用 {metrics['llm_calls']} 次 / 缓存命中 {metrics['cache_hits']} 次"
f" / 剪枝 {metrics['pruned']} 个候选（朴素做法约 {metrics['naive_calls']} 次调用）"
)

return {
"messages": [AIMessage(content="\\n".join(summary_parts))],
"metrics": metrics,
"execution_trace": [_trace_event("output", f"输出最终方案（{len(solution)} 字）")],
}

# ============================================================
# 条件边
# ============================================================
def route_after_prune(state: TreeOfThoughtState, defaults: dict | None = None) -> str:
"""剪枝后分流：有存活候选就选择，否则尝试回溯，回溯预算耗尽则收敛。"""
if state.get("beam"):
return "select"

max_backtracks = _conf(state, defaults, "max_backtracks", MAX_BACKTRACKS)
if state.get("frontier") and state.get("backtrack_count", 0) < max_backtracks:
return "backtrack"

log_warning(logger, "无存活候选且无法回溯，提前整合已有路径")
return "synthesize"

def should_continue_exploration(state: TreeOfThoughtState, defaults: dict | None = None) -> str:
"""判断是否继续下探：深度、预算、提前收敛三重边界。"""
if state.get("stop_reason"):
return "synthesize"

max_depth = _conf(state, defaults, "max_depth", MAX_DEPTH)
if state.get("current_depth", 0) >= max_depth:
log_step(logger, "Depth Limit", f"达到最大深度 {max_depth}")
return "synthesize"

if state.get("expansion_budget", EXPANSION_BUDGET) <= 0:
log_warning(logger, "展开预算耗尽，停止下探")
return "synthesize"

return "expand"

# ============================================================
# 构建 Tree of Thought Graph
# ============================================================
def build_tree_of_thought_graph(
max_depth: int = MAX_DEPTH,
max_branches: int = MAX_BRANCHES,
beam_width: int = BEAM_WIDTH,
score_threshold: float = SCORE_THRESHOLD,
satisfaction_score: float = SATISFACTION_SCORE,
max_backtracks: int = MAX_BACKTRACKS,
expansion_budget: int = EXPANSION_BUDGET,
):
"""构建多级 Tree of Thought 搜索图。

图结构（8 节点 / 3 个条件路由）：

START → decompose → expand → evaluate → prune ─┬→ select ─┬→ expand（继续下探）
│ └→ synthesize → output → END
├→ backtrack ─┬→ expand
│ └→ synthesize
└→ synthesize

这里的参数是"建图默认值"，会绑定到各节点；invoke 时 state 里传的同名字段
优先级更高，方便同一张图跑不同搜索预算。
"""
defaults = {
"max_depth": max_depth,
"max_branches": max_branches,
"beam_width": beam_width,
"score_threshold": score_threshold,
"satisfaction_score": satisfaction_score,
"max_backtracks": max_backtracks,
"expansion_budget": expansion_budget,
}

graph = StateGraph(TreeOfThoughtState)

graph.add_node("decompose", partial(decompose_node, defaults=defaults))
graph.add_node("expand", partial(expand_node, defaults=defaults))
graph.add_node("evaluate", partial(evaluate_node, defaults=defaults))
graph.add_node("prune", partial(prune_node, defaults=defaults))
graph.add_node("select", partial(select_node, defaults=defaults))
graph.add_node("backtrack", partial(backtrack_node, defaults=defaults))
graph.add_node("synthesize", partial(synthesize_node, defaults=defaults))
graph.add_node("output", partial(output_node, defaults=defaults))

graph.add_edge(START, "decompose")
graph.add_edge("decompose", "expand")
graph.add_edge("expand", "evaluate")
graph.add_edge("evaluate", "prune")
graph.add_conditional_edges(
"prune",
partial(route_after_prune, defaults=defaults),
{"select": "select", "backtrack": "backtrack", "synthesize": "synthesize"},
)
graph.add_conditional_edges(
"select",
partial(should_continue_exploration, defaults=defaults),
{"expand": "expand", "synthesize": "synthesize"},
)
graph.add_conditional_edges(
"backtrack",
partial(should_continue_exploration, defaults=defaults),
{"expand": "expand", "synthesize": "synthesize"},
)
graph.add_edge("synthesize", "output")
graph.add_edge("output", END)

return graph

def _initial_state(
task: str,
max_depth: int = MAX_DEPTH,
max_branches: int = MAX_BRANCHES,
beam_width: int = BEAM_WIDTH,
score_threshold: float = SCORE_THRESHOLD,
satisfaction_score: float = SATISFACTION_SCORE,
max_backtracks: int = MAX_BACKTRACKS,
expansion_budget: int = EXPANSION_BUDGET,
) -> dict:
"""构造初始状态。"""
return {
"messages": [HumanMessage(content=task)],
"task": task,
"branches": [],
"selected_branch": None,
"exploration_history": [],
"current_depth": 0,
"max_depth": max_depth,
"max_branches": max_branches,
"best_solution": "",
"backtrack_count": 0,
"stage_goals": [],
"criteria": [],
"path": [],
"beam": [],
"frontier": [],
"pruned": [],
"beam_width": beam_width,
"score_threshold": score_threshold,
"satisfaction_score": satisfaction_score,
"max_backtracks": max_backtracks,
"expansion_budget": expansion_budget,
"final_solution": "",
"stop_reason": "",
"metrics": dict(_DEFAULT_METRICS),
}

_RUN_CONFIG = {"recursion_limit": 60}

def _print_execution_trace(trace: list[dict]) -> None:
"""打印全程执行轨迹（多级链路的可观测性）。"""
print("\\n — 执行轨迹（多级链路） —")
for i, event in enumerate(trace, 1):
print(f" {i:02d}. [{event['ts']}] {event['node']:<11} {event['detail']}")

# ============================================================
# 运行演示
# ============================================================
def demo_architecture_design():
"""演示 1：多级链路的完整分层搜索（分解 → 发散 → 批量评估 → 剪枝 → 选择 → 整合）。"""
print("\\n— 演示 1: 架构设计的分层思维树搜索 —\\n")
_reset_stats()

graph = build_tree_of_thought_graph(max_depth=3, max_branches=3, beam_width=2)
app = graph.compile()

task = "设计一个高并发订单系统的架构方案"
print(f" 任务: {task}")
print(" max_depth=3, max_branches=3, beam_width=2")

result = app.invoke(
_initial_state(task, max_depth=3, max_branches=3, beam_width=2), _RUN_CONFIG
)

print(f"\\n {'─' * 50}")
print(f" 探索深度: {result['current_depth']} | 决策路径 {len(result.get('path', []))} 层")
_print_execution_trace(result["execution_trace"])
return result

def demo_problem_solving():
"""演示 2：剪枝撞上死胡同后回溯到 frontier 备选路径。

beam_width=2 意味着每层保留 2 条候选：第 1 名立即下探，第 2 名沉淀进备选池；
第 2 层被注入"集体违反硬约束"后，搜索不会直接失败，而是回到备选池换条路。
"""
global _DEAD_END_AT_DEPTH

print("\\n— 演示 2: 死胡同剪枝与回溯 —\\n")
_reset_stats()

graph = build_tree_of_thought_graph(max_depth=3, max_branches=3, beam_width=2)
app = graph.compile()

task = "解决生产环境的性能瓶颈问题"
print(f" 任务: {task}")
print(" beam_width=2（第 2 名进入备选池，是回溯的前提）")
print(" （死胡同注入: 第 2 层候选全部判为违反硬约束）")

_DEAD_END_AT_DEPTH = 1
try:
result = app.invoke(
_initial_state(task, max_depth=3, max_branches=3, beam_width=2), _RUN_CONFIG
)
finally:
_DEAD_END_AT_DEPTH = None

metrics = result.get("metrics", {})
print(f"\\n 探索深度: {result['current_depth']} | 回溯次数: {result['backtrack_count']}")
print(f" 剪枝候选: {metrics.get('pruned', 0)} 个 | 备选池剩余: {len(result.get('frontier', []))}")
_print_execution_trace(result["execution_trace"])
return result

def demo_depth_limit():
"""演示 3：深度与展开预算双重限制（防止组合爆炸）。"""
print("\\n— 演示 3: 深度与预算防护 —\\n")
_reset_stats()

graph = build_tree_of_thought_graph(max_depth=1, max_branches=3, expansion_budget=1)
app = graph.compile()

task = "用创新方法解决一个复杂优化问题"
print(f" 任务: {task}")
print(" max_depth=1, expansion_budget=1（只允许一层展开）")

result = app.invoke(
_initial_state(task, max_depth=1, max_branches=3, expansion_budget=1), _RUN_CONFIG
)

print(f"\\n 实际探索深度: {result['current_depth']}")
print(f" 深度限制触发: {'是' if result['current_depth'] <= 1 else '否'}")
print(f" 结束原因: {result.get('stop_reason') or '达到搜索边界'}")
return result

def demo_cache_reuse():
"""演示 4：提前收敛 + 记忆化缓存（两种省调用手段叠加）。

满意阈值放宽到 9.0：搜索够 2 层且拿到高分就不再下探第 3 层；
随后重跑同一任务，上下文签名一致，真实调用降为 0。
"""
print("\\n— 演示 4: 提前收敛与缓存复用 —\\n")

params = {"max_depth": 3, "max_branches": 3, "beam_width": 2, "satisfaction_score": 9.0}
graph = build_tree_of_thought_graph(**params)
app = graph.compile()
task = "设计一个可扩展的实时数据同步方案"

_reset_stats()
print(f" 任务: {task}")
print(" max_depth=3, satisfaction_score=9.0（够好就不下探第 3 层）")
print(" 第 1 轮：冷启动（全部为真实调用）")
first = app.invoke(_initial_state(task, **params), _RUN_CONFIG)
first_metrics = first.get("metrics", {})

_reset_stats()
print("\\n 第 2 轮：相同任务重跑（上下文签名一致，全部命中缓存）")
second = app.invoke(_initial_state(task, **params), _RUN_CONFIG)
second_metrics = second.get("metrics", {})

print(f"\\n {'─' * 50}")
print(
f" 第 1 轮: 真实调用 {first_metrics.get('llm_calls', 0)} 次"
f" / 缓存命中 {first_metrics.get('cache_hits', 0)} 次"
)
print(
f" 第 2 轮: 真实调用 {second_metrics.get('llm_calls', 0)} 次"
f" / 缓存命中 {second_metrics.get('cache_hits', 0)} 次"
)
print(f" 朴素 ToT 预计调用: {first_metrics.get('naive_calls', 0)} 次（每候选单独评估）")
print(f" 结束原因: {first.get('stop_reason') or '达到搜索边界'}")
return second

def run_demo() -> dict:
"""运行 Tree of Thought 全部演示。"""
print("=" * 60)
print(" Demo 05: Tree of Thought — 多级搜索链路与效率优化（真实 LLM）")
print("=" * 60)

arch_result = demo_architecture_design()
problem_result = demo_problem_solving()
depth_result = demo_depth_limit()
cache_result = demo_cache_reuse()

print()
print("=" * 60)
print(" 关键概念回顾")
print("=" * 60)
print(" 1. Decompose : 真实 LLM 把任务拆成逐层决策目标 + 统一评估维度")
print(" 2. Expand : 真实 LLM 承接已选路径发散候选，避开已剪枝方向")
print(" 3. Evaluate : 整层候选合并为 1 次调用批量打分（N 次 → 1 次）")
print(" 4. Prune : 束搜索 + 阈值剪枝，砍掉一个候选等于砍掉整棵子树")
print(" 5. Select : 推进深度；分数达到满意阈值即提前收敛")
print(" 6. Backtrack : 本层全灭时从 frontier 换条路，受 max_backtracks 约束")
print(" 7. Synthesize : 真实 LLM 把整条决策路径整合成可落地方案")
print(" 8. 效率优化 : 批量评估 + 束宽剪枝 + 提前收敛 + 记忆化缓存")
print()

return {
"arch_result": arch_result,
"problem_result": problem_result,
"depth_result": depth_result,
"cache_result": cache_result,
}

if __name__ == "__main__":
run_demo()

三、核心代码展开详细讲解

（一）整条链路：8 个节点、3 个条件路由

本文的核心图结构是：

START → decompose → expand → evaluate → prune ─┬→ select ─┬→ expand（继续下探）
│ └→ synthesize → output → END
├→ backtrack ─┬→ expand
│ └→ synthesize
└→ synthesize

每个节点分别做什么？

搜索一共有四个出口，只要满足一个就会收敛：

• 达到 max_depth；

• expansion_budget 耗尽；

• 当前最优分达到满意阈值，并且已经探索了足够层数；

• 当前层全灭，同时 frontier 或 max_backtracks 已经无法继续提供回溯空间。

这几个出口一起保证：模型可以探索，但不能无限探索；可以回溯，但不能无限横跳。

（二）状态设计：path / frontier / pruned 才是 ToT 真正的“记忆”

ToT 的难点不只是 prompt，而是状态怎么设计。核心状态如下，代码保持原样：

class TreeOfThoughtState(TypedDict, total=False):
messages: Annotated[list[BaseMessage], add_messages]
task: str
branches: list[ThoughtBranch]
selected_branch: ThoughtBranch | None
exploration_history: list[dict]
current_depth: int
max_depth: int
max_branches: int
best_solution: str
backtrack_count: int
# 分层搜索上下文
stage_goals: list[str] # 每一层要回答的决策问题（LLM 拆解）
criteria: list[str] # 统一评估维度（LLM 拆解）
path: list[ThoughtBranch] # 已选定的决策路径（根 → 当前）
beam: list[ThoughtBranch] # 本层剪枝后存活的候选
frontier: list[ThoughtBranch] # 全局备选池（回溯用）
pruned: list[dict] # 被剪掉的候选（供后续层避免重复思路）
# 搜索预算与阈值
beam_width: int
score_threshold: float
satisfaction_score: float
max_backtracks: int
expansion_budget: int
# 结果与可观测性
final_solution: str
stop_reason: str
metrics: dict
execution_trace: Annotated[list[dict], operator.add]

这组字段可以拆成三块理解：

1. 搜索上下文：决定“下一步从哪里继续想”

• path：当前已经选定的主路径。下一层 expand 必须看到它。

• frontier：本层没有立刻下探、但仍值得保留的备选。回溯时从这里取。

• pruned：已经被判定低价值或不可行的方向。后续发散时明确告诉模型“不要再换个说法提回来”。

2. 预算闸门：决定“最多能烧多少”

• max_depth：最多探索几层；

• max_branches：每层最多生成几条候选；

• beam_width：每层最多保留几条；

• expansion_budget：整次任务最多允许几次展开；

• max_backtracks：最多允许回溯多少次。

3. 可观测性：决定“能不能复盘与算账”

• execution_trace：记录每个节点发生了什么；

• metrics：记录真实 LLM 调用、缓存命中、生成候选、剪枝数量和朴素 ToT 预计调用；

• stop_reason：说明最终为什么结束；

• final_solution：整合后的最终方案。

对生产系统来说，这一组可观测字段不是锦上添花。没有它，你只知道“结果出来了”，却不知道模型为什么选它、为什么停、到底用了多少调用。

（三）主要的处理逻辑详细说明

1. 第一把省钱刀：批量评估，把 N 次调用压成 1 次

朴素 ToT 最直接的写法是：每个候选单独调一次模型评分。

例如一层有 3 个候选，就需要 3 次评估调用。本文把整个候选集合放进同一次模型调用里：

_EVALUATE_SYSTEM_PROMPT = """你是严格的方案评估专家，一次性评估同一层的全部候选方案。

只输出一个 JSON 对象（不要输出任何其他文字）：
{"scores": [{"id": "<候选ID>", "score": <0-10 数字>,
"dimensions": {"<维度名>": <0-10 数字>},
"is_viable": true, "comment": "<一句话理由>"}]}

要求：
1. 必须为每个候选给出一条记录，id 与输入完全一致
2. score 为综合分，评分严格并拉开差距，禁止所有候选同分
3. dimensions 使用给定的评估维度
4. 违反约束、无法落地的候选 is_viable 置为 false"""

批量评估的收益不只是“省 2 次调用”。更重要的是：模型终于有参照物了。

孤立评分时，一个方案单独拿出来，模型很容易给出“7.5、8.0、8.0”这类安全分；但把 3 个候选放到同一上下文里横向比较，模型更容易拉开差距。

原始演示里就出现了 9.0 / 8.0 / 5.0，甚至 9.0 / 6.0 / 3.0 这种更有区分度的结果。这样 score_threshold 和排序才真正有意义。

同时，LLM 的 JSON 输出仍然必须经过容错解析和字段校验。任何缺 ID、分数非法、结构不完整的记录都不能直接相信；必要时回退到启发式保底逻辑，保证流程继续运行。

2. 第二把省钱刀：剪枝——砍掉一个节点，等于砍掉整棵子树

真正把搜索空间从“指数爆炸”拉回可控范围的，是 prune。代码的核心逻辑：

survivors = (
[]
if dead_end
else [b for b in branches if b.get("is_viable", True) and b.get("score", 0) >= threshold]
)
survivors.sort(key=lambda b: b.get("score", 0), reverse=True)

# beam[0] 立即下探，beam[1:] 进备选池等回溯；束宽之外的候选直接淘汰
beam = survivors[:beam_width]
alternates = beam[1:]
dropped = [
(b, _prune_reason(b, threshold, dead_end)) for b in branches if b not in survivors
]
dropped += [(b, f"超出束宽 {beam_width}") for b in survivors[beam_width:]]

本层候选最后会被分成三类：

• beam[0]：立刻下探；

• beam[1:]：进入 frontier，作为回溯备选；

• 其余候选：彻底剪掉，不再产生下一层开销。

这就是 ToT 最关键的成本杠杆：剪枝真正省掉的，不是“这个候选本身”，而是这个候选下面原本会继续生成、评估的整棵子树。

一个很容易写错的地方：frontier 应该收谁？

这段代码里最值得记住的是：

alternates = beam[1:]

因为 beam_width=2 的正确含义应该是：

• 第 1 名现在走；

• 第 2 名先排队，未来可能回溯回来；

• 第 3 名及以后直接被束宽淘汰。

如果错误地把 survivors[beam_width:] 放进备选池，那么默认 beam_width=2 时，真正的第 2 名反而进不了 frontier。结果就是：代码里明明写了回溯逻辑，运行时却永远没有可回溯对象。

3. 第三把省钱刀：早停——“够好”不等于“第一层高分就停”

早停逻辑看起来很简单：当前最优分达到满意阈值就结束。但真实运行很容易踩坑：模型给头名打分偏高，如果一看到 9.0 就停，精心设计的多层链路会直接退化成“一层决策”。

所以我们增加了最小探索深度：

# 早停要求"分数足够高"且"已探索足够层数"：只凭第一层的高分就收敛，
# 等于放弃了分层决策本身的价值
stop_reason = ""
min_depth = min(MIN_EXPLORED_DEPTH, max_depth)
if best.get("score", 0) >= satisfaction and depth + 1 >= min_depth:
stop_reason = f"提前收敛（{best.get('score', 0)} 分 ≥ 满意阈值 {satisfaction}）"
print(f" [早停] {stop_reason}，跳过后续层的发散与评估")

这里有两个工程结论：

• 阈值不能拍脑袋设。先跑几轮，看你的模型头名分数实际分布，再定 satisfaction_score。

• 早停最好和最小探索深度一起用。这样既能省成本，又不会破坏“分层决策”的价值。

早停买的是“边际收益已经不值得继续付费”，而不是数学意义上的全局最优。

4. 第四把省钱刀：记忆化缓存——回溯与重跑不再重复付费

缓存的实现非常直接：

def _cached_invoke(key: str, messages: list[BaseMessage]) -> str:
"""带记忆化的 LLM 调用：命中缓存时不发起真实请求。"""
if key in _LLM_CACHE:
_CALL_STATS["cache_hits"] += 1
return _LLM_CACHE[key]

response = _get_tot_llm().invoke(messages)
text = str(response.content).strip()
_LLM_CACHE[key] = text
_CALL_STATS["llm_calls"] += 1
return text

但真正容易出问题的不是“有没有缓存”，而是缓存键里包含什么。本文的缓存键会带上：

• 调用类型；

• 任务；

• 层级；

• 当前层目标；

• 已选路径签名；

• 已剪枝信息等上下文。

尤其是已选路径签名不能省。

因为回溯之后，你虽然仍然在“第 2 层”，但上层选择已经变了。如果只按“任务 + 第几层”缓存，模型会直接拿回上一条路生成过的候选，相当于回溯后继续重复旧答案。

换句话说：同一层 ≠ 同一上下文。缓存必须绑定搜索路径。

（四）实战代码：图结构本身反而是最简单的部分

8 个节点、3 个条件路由在 LangGraph 里非常直白，代码如下：

graph.add_edge(START, "decompose")
graph.add_edge("decompose", "expand")
graph.add_edge("expand", "evaluate")
graph.add_edge("evaluate", "prune")
graph.add_conditional_edges(
"prune",
partial(route_after_prune, defaults=defaults),
{"select": "select", "backtrack": "backtrack", "synthesize": "synthesize"},
)
graph.add_conditional_edges(
"select",
partial(should_continue_exploration, defaults=defaults),
{"expand": "expand", "synthesize": "synthesize"},
)
graph.add_conditional_edges(
"backtrack",
partial(should_continue_exploration, defaults=defaults),
{"expand": "expand", "synthesize": "synthesize"},
)
graph.add_edge("synthesize", "output")
graph.add_edge("output", END)

select 和 backtrack 共用同一套“还要不要继续探索”的判断，因为无论是正常推进，还是刚刚换了一条路线，下一步都只关心三件事：

• 是否已经触发早停；

• 是否达到最大深度；

• 是否还有展开预算。

四、运行实记与演示说明

（一）运行方式

source .venv/bin/activate
python stages/stage4_reasoning/05_tree_of_thought/main.py

另外，多级链路每层会经过多个节点，超步数量远高于线性 Agent。代码把 recursion_limit 提到 60，就是为了避免多层搜索轻易撞上默认限制。

（二）具体演示结果详细说明

1. 运行演示 1：三层决策真的串起来了

第一组演示配置：

max_depth=3, max_branches=3, beam_width=2
任务：设计一个高并发订单系统的架构方案

第一步：先由 LLM 拆成 3 个逐层目标

第 1 层目标: 确定系统的核心架构模式：单体、微服务还是事件驱动架构
第 2 层目标: 在选定架构下，选择订单处理的并发控制机制：队列异步处理、分布式锁还是乐观并发控制
第 3 层目标: 设计数据存储与一致性方案：数据库选型、分库分表策略、最终一致性实现方式
评估维度: 可行性 / 性能 / 成本 / 风险

注意第 2 层目标里的那句：“在选定架构下”。这说明第二层并不是重新思考“订单系统怎么设计”，而是只回答在第一层决策成立之后才出现的问题。

第二步：第一层批量评估

第一层的真实结果是：

于是：

• 事件驱动架构：立即下探；

• 微服务架构：进入 frontier；

• 单体架构：超出束宽，被剪枝。

第三步：第二层候选开始“带着上一层答案思考”

真实生成的候选包括：

[d1_a] 基于事件序号的有序异步队列处理
[d1_b] 基于版本号的分布式乐观并发控制
[d1_c] 基于资源分片的局部分布式锁

最有说服力的地方就在这里：这些候选已经明显围绕事件驱动架构继续细化，而不是重新讨论“单体还是微服务”。

第三层最终选到“按用户 ID 哈希分库分表 + 跨库事务补偿表”，三层决策被整合成一个完整方案。

这次搜索一共花了多少？

真实调用：8 次
缓存命中：0 次
剪枝：3 个候选
朴素做法预计：14 次调用

8 次调用的构成非常清楚：

• 1 次 decompose；

• 3 层 ×（1 次 expand + 1 次批量 evaluate）= 6 次；

• 1 次 synthesize。

也就是：3 层完整决策只用了 8 次真实调用，朴素做法约需 14 次，减少约 43%。

省下来的 6 次，主要来自“每层 3 个候选不再分别调用评估模型”。

2. 运行演示 2：撞上死胡同，ToT 才真正体现价值

第二组演示专门注入一个死胡同：第 2 层所有候选都被判定违反硬约束。第 1 层先保留两条路线：

第 1 名：数据库慢查询日志 + 执行计划分析（9.0）→ 立即下探
第 2 名：分布式追踪分析应用调用链（8.0）→ 进入 frontier

第一次进入第 2 层后，3 个候选全部被剪掉：

[剪枝] 第 2 层候选全部违反硬约束
→ 本层无存活候选

正常的线性推理到这里往往只能“失败”或“硬着头皮选一个”。ToT 则会：

• 发现 beam 为空；

• 检查 frontier 是否还有候选；

• 从全局备选池拿到第 1 层第二名；

• 把 path 截断到对应深度；

• 接上新的上层选择，重新进入第 2 层展开。

回溯前后，第 2 层的候选发生了明显变化：

回溯前：围绕数据库 / JVM / 系统性能剖析去查瓶颈。

回溯后：变成方法级插桩、上下文聚类、异步队列反推等更贴合“分布式追踪”路线的方案。

这证明两件事：

• 上层 path 确实传进了下一层；

• 缓存键必须带 path，否则回溯会命中旧缓存，得到完全相同的候选。

回溯不是免费的

这次完整搜索：

真实调用：10 次
朴素做法预计：18 次
回溯次数：1
剪枝候选：6 个

相比演示 1，多出来的 2 次真实调用正是：

• 回溯后的第 2 层重新 expand；

• 重新 evaluate。

所以 max_backtracks 必须存在。允许换路，不代表允许无限横跳。

3. 运行演示 3：深度 + 展开预算，给搜索空间上“双保险”

第三组演示把限制压得非常紧：

max_depth=1
expansion_budget=1

只允许一层展开。真实评估中，有一个“量子退火”方案拿到了全场最高的创新性 10 分，但可行性只有 2 分，并且被判定 is_viable=false，直接淘汰。

这非常适合说明一个工程原则：

“最创新”与“最可落地”往往不是同一个方案。可行性应当拥有一票否决能力，而不是只参与加权平均。

如果没有预算闸门，假设每层 3 个候选、探索 5 层，理论路径数会达到 3^5 = 243。

因此至少要同时控制：

• max_depth：锁深度；

• max_branches：锁每层生成数量；

• beam_width：锁每层真正存活数量；

• expansion_budget：锁总展开次数。

这四把锁合在一起，才真正把组合爆炸关在门外。

4. 运行演示 4：早停 + 缓存，两种省法可以叠加

第四组演示把满意阈值调整到 9.0，最大深度仍是 3。

任务：

设计一个可扩展的实时数据同步方案

第一轮：

第 1 层选定：基于发布订阅模式的事件驱动同步（9.0）
第 2 层选定：基于 Kafka 的分区有序 + 幂等消费实现最终一致（9.0）
→ 满足早停条件，不再展开第 3 层

因此原本规划好的“节点扩展与故障恢复机制”这一层没有继续发散，直接省下后续两次调用。

第一轮统计：

真实调用：6 次
缓存命中：0 次
朴素做法预计：10 次

紧接着，用完全相同的任务和上下文再跑一遍：

真实调用：0 次
缓存命中：6 次

这就是记忆化缓存最直观的收益：调 prompt、调阈值、重复测试同一输入时，如果上下文签名没有变化，就不应该一遍遍重新付费。

（三）把账算到图上：优化后的调用次数到底省了多少？

工程里真正有价值的不是一句“这个方案更省 Token”，而是能够直接回答：

• 真实发了几次请求？

• 命中几次缓存？

• 生成多少候选？

• 剪掉多少候选？

• 如果用朴素方案，大概要发多少次请求？

• 本次为什么结束？

这也是为什么 metrics 和 stop_reason 值得成为正式状态，而不是临时 print。

五、问题排查与线上注意事项

（一）六个最容易踩的坑，以及应该怎么排查

坑 1：早停阈值太低，多级链路第一层就结束

现象：设计了 3 层，实际每次只跑 1 层。

原因：模型评分偏高，头名经常直接到 9.0。

处理：先跑几轮观察分数分布，再校准 satisfaction_score；同时增加最小探索深度。

坑 2：评估维度名对不上，维度分悄悄丢失

现象：某些演示能打印完整维度分，某些只剩综合分。

原因：decompose 可能生成“可行性（实施难度和资源需求）”这类长名字，但 evaluate 返回的 key 变成“可行性”，按字符串精确匹配时取不到。

处理：

• 在分解 prompt 中限制维度名短一些，例如 2～6 个字、不带括号；或

• 在评分索引阶段做名称归一化。

这种问题危险在于它通常不抛异常，只是静默退化。

坑 3：备选池一直为空，回溯逻辑实际上是死代码

现象：写了 backtrack，运行却从没触发过。

原因：把“束宽之外”的候选放进 frontier，而不是把 beam[1:] 放进去。

处理：明确语义——beam 表示“本层仍允许参与后续搜索的候选”，其中第一名现在走，其余才是未来回溯对象。

坑 4：缓存键漏掉 path，回溯之后生成一模一样的候选

现象：明明换了上层路线，第 2 层却和之前完全一样。

原因：缓存键只包含“任务 + 层级”，没有包含已选路径。

处理：所有“同一层但上下文可能不同”的调用，都必须把完整路径签名放进 key。

坑 5：一调大深度和分支数，账单和运行时间一起爆炸

现象：运行越来越慢，还可能撞 GraphRecursionError。

原因：搜索空间天然是指数关系，而且多级链路每层还要经过多个图节点。

处理：四把锁一起用——max_depth / max_branches / beam_width / expansion_budget；再根据层数和回溯余量显式放大 recursion_limit。

坑 6：每个候选单独打分，结果全部 8 分上下

现象：看起来每个方案都“挺好”，根本拉不开差距。

原因：孤立评分缺少参照，模型容易给安全分。

处理：整层候选一次输入，明确要求横向比较并“禁止所有候选同分”。

（二）生产级还应该补什么？

本文的实现已经具备完整的搜索、剪枝、回溯和成本控制骨架，但如果真要进入生产，还可以继续增强：

1. 并行展开

当前 beam_width=2 时，仍然只让当前最优的一条立即下探。生产环境可以对 beam 内多条路径使用 llm.batch 或异步并发，同时探索，用更多计算换更高搜索质量。

2. 跨会话共享缓存

把进程内 _LLM_CACHE 换成 Redis 一类共享存储。团队调 prompt 或重复处理相同上下文时，可以跨进程复用结果。

3. frontier 持久化

把备选池写进 Checkpointer，让长任务可以中断恢复，甚至在约束变化后从过去的备选路径继续搜索。

4. 发散模型和评估模型分开

发散阶段更需要创造力，评估阶段更需要稳定、严格和一致。把两种角色交给不同模型，可以避免“同一个模型既出题又给自己高分”。

5. 和其他推理模式组合

• ToT + Reflection：先搜索出当前最优路径，再对最终方案做反思改进；

• ToT + Self-Correction：先搜索方案，再执行验证；失败后回溯到备选池换方案。

（三）什么时候值得用 ToT，什么时候不值得？

ToT 不应该成为所有 Agent 的默认推理方式。

适合使用

• 方案空间明显有多个合理候选；

• 前一层决策会改变后一层问题；

• 选错路线的代价较高；

• 有明确的可比较维度；

• 允许多花一些 Token 换更好的决策质量；

• 需要保留“备选路线”，而不是只要一个答案。

例如：架构选型、复杂故障定位、策略规划、约束优化、多阶段设计。

不太适合

• 问题有明确唯一答案；

• 任务极简单，单次推理足够；

• 延迟要求非常苛刻；

• 成本预算极低；

• 各候选之间根本没有稳定的评估标准。

对这些任务，上 ToT 只会把一个简单问题变成昂贵流程。

六、Stage 4 收官：真正能上线的推理 Agent，都必须有“硬边界”

到这里，Stage 4 的五种推理模式就完整了：

• ReAct：边想边做，工具增强；

• Planning：先规划，再分步执行；

• Reflection：生成后自我审视、迭代提升；

• Self-Correction：执行后验证、自动纠错；

• Tree of Thought：分层发散、评估剪枝、选优回溯。

它们看起来各不相同，但有一条共同的工程主线：

任何能循环、重试、反思、搜索的 Agent，都必须有防死循环、防组合爆炸、防无限付费的硬性保护。

不同模式对应不同闸门：

• max_iterations

• max_steps

• max_rounds

• max_retries

• max_depth

• expansion_budget

• max_backtracks

• satisfaction / early stop

这才是把“炫技 Demo”变成“敢放进生产环境”的底线。

Tree of Thought 的价值也可以用一句话收尾：

它不是替你把所有路都走一遍，而是让 AI 知道什么时候应该多想、什么时候应该放弃、什么时候值得换路，以及什么时候已经足够好。

而且这四件事，都必须能被预算、状态和指标约束。

下一阶段，问题会从“一个 Agent 怎么想得更好”升级成“多个专精 Agent 怎么协同工作”。

👉 下一篇（Stage 5 开篇）：《从单打独斗到团队协作：用 LangGraph 构建 Supervisor-Worker 多 Agent 系统》。

系列导航：LangGraph从零构建生产级 AI Agent 平台的递进式学习项目-CSDN博客

附：本文核心参数速查

最后记住四个词就够了：批量、剪枝、早停、缓存。

它们分别解决：本层评估太贵、子树太多、层数太深、重复调用太浪费。