Chain-of-Thought 进阶：Tree-of-Thought 与 Graph-of-Thought

问题 → [思考步骤1] → [思考步骤2] → [思考步骤3] → 答案

请一步步思考：
1. 首先分析问题中的关键条件...
2. 然后推导中间结论...
3. 最后得出最终答案...

                    [初始状态]
                   /          \
            [思路A]           [思路B]
           /        \         /        \
      [A-1]       [A-2]   [B-1]       [B-2]
        |           |        |
      [答案A]    [答案A']  [答案B]
      
      评估：答案B 最优 → 选择路径 B → B-1

from typing import List, Optional, Callable
from dataclasses import dataclass, field
import json

@dataclass
class ThoughtNode:
    """思维树节点"""
    state: str           # 当前思维状态描述
    thought: str         # 到达此状态的思考内容
    parent: Optional['ThoughtNode'] = None
    children: List['ThoughtNode'] = field(default_factory=list)
    value: float = 0.0   # 评估值 0-1
    depth: int = 0
    visited: bool = False

class TreeOfThought:
    """Tree-of-Thought 推理引擎"""
    
    def __init__(self, llm_client, max_depth: int = 5, 
                 branching_factor: int = 3, beam_size: int = 2):
        self.llm = llm_client
        self.max_depth = max_depth
        self.branching = branching_factor
        self.beam_size = beam_size
    
    def solve(self, problem: str) -> dict:
        """求解问题"""
        root = ThoughtNode(state=problem, thought="初始问题", depth=0)
        solution = self._search(root)
        return {
            'answer': solution.thought if solution else None,
            'path': self._trace_path(solution) if solution else [],
            'tree_stats': {
                'nodes_generated': self._count_nodes(root),
                'max_depth_reached': self._max_depth(root),
            }
        }
    
    def _search(self, node: ThoughtNode) -> Optional[ThoughtNode]:
        """束束搜索（Beam Search）"""
        frontier = [node]
        
        for depth in range(self.max_depth):
            next_frontier = []
            
            for current in frontier:
                if self._is_solution(current):
                    return current
                
                # 生成多个思维分支
                thoughts = self._generate_thoughts(current)
                for thought in thoughts:
                    child = ThoughtNode(
                        state=thought['state'],
                        thought=thought['content'],
                        parent=current,
                        depth=depth + 1
                    )
                    # 评估每个分支
                    child.value = self._evaluate(child)
                    current.children.append(child)
                    next_frontier.append(child)
            
            # 保留 top-k 分支
            next_frontier.sort(key=lambda n: -n.value)
            frontier = next_frontier[:self.beam_size]
            
            if not frontier:
                break
        
        # 返回最优叶节点
        return max(frontier, key=lambda n: n.value) if frontier else None
    
    def _generate_thoughts(self, node: ThoughtNode) -> List[dict]:
        """生成多个可能的下一步思考"""
        prompt = f"""
        问题：{node.state}
        当前思考：{node.thought}
        当前深度：{node.depth}
        
        请生成 {self.branching} 个不同的下一步思考方向。
        每个方向应探索不同的推理路径。
        
        输出JSON格式：
        [
            {{"state": "更新后的问题状态", "content": "具体思考内容"}},
            ...
        ]
        """
        response = self.llm.generate(prompt)
        return json.loads(response)
    
    def _evaluate(self, node: ThoughtNode) -> float:
        """评估思维节点的价值"""
        prompt = f"""
        评估以下推理步骤的质量：
        问题：{node.parent.state if node.parent else node.state}
        推理步骤：{node.thought}
        
        请从以下维度评分（0-1）：
        1. 逻辑正确性
        2. 与问题相关性
        3. 推进进度（离答案有多近）
        
        返回平均分。
        """
        response = self.llm.generate(prompt)
        return float(response.strip())
    
    def _is_solution(self, node: ThoughtNode) -> bool:
        """判断是否已到达答案"""
        prompt = f"以下内容是否已经给出了问题的完整答案？回答是或否。\n{node.thought}"
        return "是" in self.llm.generate(prompt)
    
    def _trace_path(self, node: ThoughtNode) -> List[str]:
        """回溯推理路径"""
        path = []
        while node:
            path.append({'depth': node.depth, 'thought': node.thought, 'value': node.value})
            node = node.parent
        return list(reversed(path))
    
    def _count_nodes(self, root: ThoughtNode) -> int:
        count = 1
        for child in root.children:
            count += self._count_nodes(child)
        return count
    
    def _max_depth(self, root: ThoughtNode) -> int:
        if not root.children:
            return root.depth
        return max(self._max_depth(c) for c in root.children)

     [A] ──→ [B] ──→ [D] ──→ [F]
      │              ↗
      └──→ [C] ──→ [E] ──→ [G] ──→ [答案]
                     ↑
              [H] ──┘

from typing import Dict, Set, List
from collections import defaultdict

class GraphOfThought:
    """Graph-of-Thought 推理引擎"""
    
    def __init__(self, llm_client):
        self.llm = llm_client
        self.nodes: Dict[str, dict] = {}  # id → node data
        self.edges: Dict[str, Set[str]] = defaultdict(set)  # src → {dsts}
        self.reverse_edges: Dict[str, Set[str]] = defaultdict(set)  # dst → {srcs}
    
    def add_thought(self, thought_id: str, content: str, 
                    parents: List[str] = None) -> dict:
        """添加思维节点"""
        self.nodes[thought_id] = {
            'id': thought_id,
            'content': content,
            'value': 0.0,
            'aggregated': False
        }
        if parents:
            for parent_id in parents:
                self.edges[parent_id].add(thought_id)
                self.reverse_edges[thought_id].add(parent_id)
        return self.nodes[thought_id]
    
    def aggregate(self, node_ids: List[str], new_id: str) -> dict:
        """聚合多个思维节点——GoT 的核心操作"""
        contents = [self.nodes[nid]['content'] for nid in node_ids]
        
        prompt = f"""
        以下是多个推理路径的中间结论：
        {json.dumps(contents, ensure_ascii=False, indent=2)}
        
        请综合这些结论，生成一个更全面、更准确的综合判断。
        """
        merged_content = self.llm.generate(prompt)
        
        return self.add_thought(new_id, merged_content, parents=node_ids)
    
    def refine(self, node_id: str, feedback: str = "") -> dict:
        """精化思维节点"""
        node = self.nodes[node_id]
        prompt = f"""
        原始推理：{node['content']}
        反馈：{feedback}
        
        请改进以上推理，使其更准确、更完整。
        """
        improved = self.llm.generate(prompt)
        refined_id = f"{node_id}_refined"
        return self.add_thought(refined_id, improved, parents=[node_id])
    
    def solve(self, problem: str) -> dict:
        """完整 GoT 推理流程"""
        # Phase 1: 分解问题
        sub_problems = self._decompose(problem)
        
        # Phase 2: 并行推理
        for i, sub in enumerate(sub_problems):
            self.add_thought(f"sub_{i}", sub)
            # 每个子问题生成多个解法
            for j in range(3):
                solution = self._solve_sub(sub)
                self.add_thought(f"sol_{i}_{j}", solution, parents=[f"sub_{i}"])
        
        # Phase 3: 聚合
        for i in range(len(sub_problems)):
            sols = [f"sol_{i}_{j}" for j in range(3)]
            self.aggregate(sols, f"agg_{i}")
        
        # Phase 4: 整合
        all_aggs = [f"agg_{i}" for i in range(len(sub_problems))]
        final = self.aggregate(all_aggs, "final")
        
        return {
            'answer': final['content'],
            'graph_size': len(self.nodes),
            'edge_count': sum(len(v) for v in self.edges.values()),
            'sub_problems': len(sub_problems)
        }
    
    def _decompose(self, problem: str) -> List[str]:
        prompt = f"将以下问题分解为独立的子问题：\n{problem}"
        return json.loads(self.llm.generate(prompt))
    
    def _solve_sub(self, sub_problem: str) -> str:
        prompt = f"解决以下子问题：\n{sub_problem}"
        return self.llm.generate(prompt)
    
    def visualize(self) -> str:
        """生成 Graphviz 可视化"""
        dot = ["digraph GoT {"]
        dot.append('  rankdir=TB;')
        dot.append('  node [shape=box, style=filled, fillcolor=lightblue];')
        for nid, node in self.nodes.items():
            label = node['content'][:30].replace('"', "'")
            dot.append(f'  "{nid}" [label="{label}"];')
        for src, dsts in self.edges.items():
            for dst in dsts:
                dot.append(f'  "{src}" -> "{dst}";')
        dot.append("}")
        return "\n".join(dot)

def select_reasoning_method(problem: dict) -> str:
    """根据问题特征选择推理方法"""
    if problem['complexity'] == 'low':
        return 'direct'  # 直接回答
    elif problem['type'] in ['math', 'logic'] and problem['steps'] < 5:
        return 'cot'     # Chain-of-Thought
    elif problem['type'] in ['planning', 'search', 'puzzle']:
        return 'tot'     # Tree-of-Thought
    elif problem['type'] in ['proof', 'analysis', 'synthesis']:
        return 'got'     # Graph-of-Thought
    elif problem['requires_backtracking']:
        return 'tot' if problem['branches'] < 10 else 'got'
    else:
        return 'cot'