Everything of Thoughts (XoT)

UCB1(node) = Q(node) / N(node) + C * sqrt(ln(N(parent)) / N(node))

import math
import random
from dataclasses import dataclass, field
from openai import OpenAI

client = OpenAI()


@dataclass
class MCTSNode:
    """Узел дерева MCTS."""
    state: str                           # текущее состояние рассуждения
    parent: "MCTSNode | None" = None
    children: list["MCTSNode"] = field(default_factory=list)
    visits: int = 0
    total_reward: float = 0.0
    is_terminal: bool = False

    @property
    def avg_reward(self) -> float:
        return self.total_reward / self.visits if self.visits > 0 else 0.0

    def ucb1(self, exploration_constant: float = 1.41) -> float:
        """Вычислить UCB1 для этого узла."""
        if self.visits == 0:
            return float("inf")  # неисследованные узлы — приоритет
        parent_visits = self.parent.visits if self.parent else 1
        exploitation = self.avg_reward
        exploration = exploration_constant * math.sqrt(
            math.log(parent_visits) / self.visits
        )
        return exploitation + exploration


def llm_expand(node: MCTSNode, task: str, n: int = 3) -> list[MCTSNode]:
    """Expansion: LLM генерирует n вариантов следующего шага."""
    resp = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "Ты генерируешь варианты следующего шага "
                    "рассуждения. Каждый вариант на отдельной строке, "
                    "начиная с номера (1., 2., 3.). Только варианты, "
                    "без пояснений."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Задача: {task}\n\n"
                    f"Текущее состояние рассуждения:\n{node.state}\n\n"
                    f"Предложи {n} различных варианта "
                    "следующего шага."
                ),
            },
        ],
        temperature=0.8,
    )
    text = resp.choices[0].message.content
    children = []
    for line in text.strip().split("\n"):
        line = line.strip()
        if line and line[0].isdigit():
            thought = line.lstrip("0123456789.)- ").strip()
            if thought:
                child = MCTSNode(
                    state=f"{node.state}\n→ {thought}",
                    parent=node,
                )
                children.append(child)
    node.children = children[:n]
    return node.children


def llm_evaluate(state: str, task: str) -> float:
    """Simulation: LLM оценивает текущее состояние (0.0 — 1.0)."""
    resp = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": (
                    f"Задача: {task}\n\n"
                    f"Текущее рассуждение:\n{state}\n\n"
                    "Оцени, насколько это рассуждение близко к "
                    "правильному решению. Ответь ОДНИМ числом "
                    "от 0.0 до 1.0, где 1.0 — задача решена верно, "
                    "0.0 — полный тупик."
                ),
            }
        ],
        temperature=0,
    )
    text = resp.choices[0].message.content.strip()
    try:
        score = float(text.split()[0])
        return max(0.0, min(1.0, score))
    except (ValueError, IndexError):
        return 0.5


def select(node: MCTSNode) -> MCTSNode:
    """Selection: спуск по дереву через UCB1."""
    current = node
    while current.children:
        current = max(current.children, key=lambda c: c.ucb1())
    return current


def backpropagate(node: MCTSNode, reward: float) -> None:
    """Backpropagation: обновление статистик вверх по дереву."""
    current = node
    while current is not None:
        current.visits += 1
        current.total_reward += reward
        current = current.parent


def xot_solve(
    task: str,
    max_iterations: int = 20,
    expand_count: int = 3,
    reward_threshold: float = 0.9,
) -> str:
    """
    Основной цикл XoT: MCTS + LLM.

    Args:
        task: описание задачи
        max_iterations: макс. итераций MCTS
        expand_count: вариантов на каждом expansion
        reward_threshold: порог для досрочной остановки
    """
    root = MCTSNode(state="Начало рассуждения.")

    best_node = root
    best_reward = 0.0

    for i in range(max_iterations):
        # 1. Selection — выбираем лист через UCB1
        leaf = select(root)

        # 2. Expansion — LLM генерирует варианты
        if not leaf.is_terminal:
            children = llm_expand(leaf, task, n=expand_count)
            if children:
                leaf = random.choice(children)

        # 3. Simulation — LLM оценивает состояние
        reward = llm_evaluate(leaf.state, task)

        # 4. Backpropagation — обновляем статистики
        backpropagate(leaf, reward)

        # Отслеживаем лучший результат
        if reward > best_reward:
            best_reward = reward
            best_node = leaf

        # Досрочная остановка, если нашли хорошее решение
        if reward >= reward_threshold:
            break

    return best_node.state


# Пример использования
result = xot_solve(
    task="Из чисел 2, 5, 7, 11 получи ровно 24 (+, -, *, /). "
         "Каждое число — один раз.",
    max_iterations=15,
    expand_count=3,
)
print(result)