1. John 原有 5 个，给了 Mary 2 个，John 剩 5-2=3 个  ✓
2. Mary 收到 2 个，又给了 Tom 1 个，Tom 得到 1 个  ✓
3. Tom 原来有 3 个，加上 1 个，Tom 有 3+1=4 个  ✓
4. John 有 3 个，Tom 有 4 个，合计 3+4=7 个  ✓

import os
import re
from collections import Counter
from typing import Optional
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(
    model="deepseek-chat",
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com",
    # temperature 设高一点，让每次采样路径不同
    # 如果 temperature=0，每次结果一样，投票没有意义
    temperature=0.7,
)

cot_template = ChatPromptTemplate.from_messages([
    ("system", "你是一个精确的推理助手。请一步一步分析，最后给出答案。答案格式为：最终答案：[数字或结论]"),
    ("human", "{question}"),
])

def extract_answer(response: str) -> Optional[str]:
    """从推理结果中提取最终答案"""
    # 匹配"最终答案：xxx"或"答案是xxx"等模式
    patterns = [
        r"最终答案[：:]\s*(.+?)(?:\n|$)",
        r"答案[是为][：:]?\s*(.+?)(?:\n|$)",
        r"所以.{0,5}[是为有]\s*(.+?)(?:\n|$)",
    ]
    for pattern in patterns:
        match = re.search(pattern, response)
        if match:
            return match.group(1).strip()
    # 如果没有匹配到标准格式，取最后一行作为答案
    lines = [l.strip() for l in response.strip().split("\n") if l.strip()]
    return lines[-1] if lines else None

def self_consistency(question: str, num_samples: int = 5) -> dict:
    """
    Self-Consistency 推理：多次采样，取众数答案
    num_samples=5 是实践中常用的数量，再多收益递减但成本线性增加
    """
    answers = []
    reasoning_paths = []

    for i in range(num_samples):
        messages = cot_template.format_messages(question=question)
        response = llm.invoke(messages).content
        answer = extract_answer(response)

        reasoning_paths.append(response)
        if answer:
            answers.append(answer)

    # 统计各答案出现次数
    answer_counts = Counter(answers)
    majority_answer, majority_count = answer_counts.most_common(1)[0]

    return {
        "question": question,
        "majority_answer": majority_answer,
        "confidence": majority_count / num_samples,
        "answer_distribution": dict(answer_counts),
        "reasoning_paths": reasoning_paths,
        "num_samples": num_samples,
    }


# 验证目的：观察多次采样的答案分布，验证多数投票是否能纠正偶发错误
question = """
小明有24块糖，他给了小红1/3，小红又把自己的糖的一半给了小李。
如果小李原来有8块糖，现在小李有多少块糖？
"""

result = self_consistency(question, num_samples=5)
print(f"问题：{result['question'].strip()}")
print(f"最终答案：{result['majority_answer']}")
print(f"置信度：{result['confidence']:.0%}（{result['num_samples']} 次采样中的多数）")
print(f"答案分布：{result['answer_distribution']}")

import os
from dataclasses import dataclass, field
from typing import List, Optional
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

# 两个模型：一个生成想法，一个评估想法
# 实际场景中可以用同一个模型，也可以用强弱两个模型
generator_llm = ChatOpenAI(
    model="deepseek-chat",
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com",
    temperature=0.8,  # 生成时温度高，鼓励多样性
)

evaluator_llm = ChatOpenAI(
    model="deepseek-chat",
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com",
    temperature=0,  # 评估时温度低，要求一致稳定
)

@dataclass
class ThoughtNode:
    content: str           # 这个思维节点的内容
    score: float = 0.0     # 评估得分（0-10）
    depth: int = 0         # 在树中的深度
    children: List["ThoughtNode"] = field(default_factory=list)
    parent: Optional["ThoughtNode"] = None

# 生成候选思路的 Prompt
generate_prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一个严谨的逻辑推理助手。
给定一个问题和当前的推理状态，生成 {n_thoughts} 个不同的下一步思路。
每个思路独立一行，用"思路N："开头。思路要具体可执行，不要泛泛而谈。"""),
    ("human", """问题：{problem}

当前推理状态：
{current_state}

请生成 {n_thoughts} 个不同的下一步思路："""),
])

# 评估思路质量的 Prompt
evaluate_prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一个严格的推理质量评估者。
评估给定思路是否是解决问题的有效步骤。
只输出一个 0-10 的整数分数（10 分最好），不要其他内容。"""),
    ("human", """问题：{problem}

当前思路：{thought}

评分（0-10）："""),
])

def generate_thoughts(problem: str, current_state: str, n_thoughts: int = 3) -> List[str]:
    """生成多个候选思路"""
    messages = generate_prompt.format_messages(
        problem=problem,
        current_state=current_state,
        n_thoughts=n_thoughts,
    )
    response = generator_llm.invoke(messages).content
    # 解析"思路N："格式
    thoughts = []
    for line in response.split("\n"):
        line = line.strip()
        if line and ("思路" in line or line[0].isdigit()):
            # 去掉"思路N："前缀
            thought = line.split("：", 1)[-1].strip() if "：" in line else line
            if thought:
                thoughts.append(thought)
    return thoughts[:n_thoughts]

def evaluate_thought(problem: str, thought: str) -> float:
    """让 LLM 对一个思路打分"""
    messages = evaluate_prompt.format_messages(
        problem=problem,
        thought=thought,
    )
    response = evaluator_llm.invoke(messages).content.strip()
    try:
        score = float(re.search(r"\d+", response).group())
        return min(10.0, max(0.0, score))
    except (AttributeError, ValueError):
        return 5.0  # 解析失败时给中间分

def tot_solve(problem: str, max_depth: int = 3, n_thoughts: int = 3, beam_width: int = 2) -> str:
    """
    Tree of Thoughts 求解
    beam_width：每层保留的最优节点数（Beam Search）
    max_depth：最大探索深度
    """
    import re

    # 初始状态：空
    root = ThoughtNode(content="开始分析", depth=0)
    current_level = [root]

    for depth in range(1, max_depth + 1):
        next_level_candidates = []

        for node in current_level:
            # 从当前节点生成子思路
            current_state = f"已有推理：{node.content}" if node.content != "开始分析" else "尚未开始"
            thoughts = generate_thoughts(problem, current_state, n_thoughts)

            for thought_text in thoughts:
                score = evaluate_thought(problem, thought_text)
                child = ThoughtNode(
                    content=thought_text,
                    score=score,
                    depth=depth,
                    parent=node,
                )
                node.children.append(child)
                next_level_candidates.append(child)

        if not next_level_candidates:
            break

        # Beam Search：只保留得分最高的 beam_width 个节点继续探索
        next_level_candidates.sort(key=lambda x: x.score, reverse=True)
        current_level = next_level_candidates[:beam_width]

        print(f"第 {depth} 层，保留 {len(current_level)} 个节点：")
        for node in current_level:
            print(f"  [分数 {node.score:.1f}] {node.content[:60]}...")

    # 找到得分最高的叶节点，回溯整条推理链
    best_node = max(current_level, key=lambda x: x.score)

    # 回溯路径
    path = []
    node = best_node
    while node is not None and node.content != "开始分析":
        path.append(node.content)
        node = node.parent
    path.reverse()

    return "\n".join([f"步骤 {i+1}：{step}" for i, step in enumerate(path)])


# 测试
problem = "一个公司有 100 名员工，其中 60% 会编程，40% 会设计，25% 两者都会。只会编程不会设计的有多少人？"

print(f"问题：{problem}\n")
print("ToT 推理过程：")
solution = tot_solve(problem, max_depth=2, n_thoughts=3, beam_width=2)
print(f"\n最优推理路径：\n{solution}")