Buffer of Thoughts (BoT)

from dataclasses import dataclass, field

@dataclass
class ThoughtTemplate:
    """Шаблон мышления для мета-буфера."""
    id: str
    task_type: str                   # тип задачи: "arithmetic_puzzle", "geometry", ...
    strategy: str                    # стратегия: "pair_decomposition", "backward_reasoning"
    steps: list[str]                 # последовательность шагов
    examples: list[str] = field(default_factory=list)  # примеры применения
    success_count: int = 0           # сколько раз успешно применён
    fail_count: int = 0              # сколько раз не сработал

    @property
    def success_rate(self) -> float:
        total = self.success_count + self.fail_count
        return self.success_count / total if total > 0 else 0.0


class MetaBuffer:
    """Мета-буфер: хранение и поиск thought-templates."""

    def __init__(self):
        self.templates: dict[str, ThoughtTemplate] = {}

    def add(self, template: ThoughtTemplate) -> None:
        self.templates[template.id] = template

    def retrieve(self, task_type: str) -> ThoughtTemplate | None:
        """Найти лучший шаблон для данного типа задачи."""
        candidates = [
            t for t in self.templates.values()
            if t.task_type == task_type
        ]
        if not candidates:
            return None
        # Выбираем шаблон с лучшим success_rate
        return max(candidates, key=lambda t: (t.success_rate, t.success_count))

    def update_stats(self, template_id: str, success: bool) -> None:
        """Обновить статистику после использования."""
        if template_id in self.templates:
            t = self.templates[template_id]
            if success:
                t.success_count += 1
            else:
                t.fail_count += 1

from openai import OpenAI

client = OpenAI()

CLASSIFY_PROMPT = """Определи тип задачи. Варианты:
- arithmetic_puzzle (числовые головоломки, Game of 24)
- geometry (геометрические фигуры, координаты)
- chess (шахматные задачи, мат в N ходов)
- logic (логические задачи, выводы)
- math_proof (математические доказательства)
- text_analysis (анализ текста, классификация)

Задача: {task}

Ответь ОДНИМ словом — типом задачи."""

INSTANTIATE_PROMPT = """У тебя есть шаблон мышления:
Тип: {task_type}
Стратегия: {strategy}
Шаги:
{steps}

Адаптируй шаблон к конкретной задаче:
{task}

Запиши адаптированный план решения."""

REASON_PROMPT = """Реши задачу строго по плану.

Задача: {task}

План решения:
{plan}

Выполни каждый шаг и дай финальный ответ."""


def buffer_of_thoughts(
    task: str,
    buffer: MetaBuffer,
    model: str = "gpt-4o"
) -> dict:
    """Полный цикл Buffer of Thoughts."""

    # Шаг 1: Classify — определяем тип задачи
    classify_resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": CLASSIFY_PROMPT.format(task=task)}],
        temperature=0
    )
    task_type = classify_resp.choices[0].message.content.strip().lower()

    # Шаг 2: Retrieve — ищем шаблон
    template = buffer.retrieve(task_type)

    if template is None:
        # Нет шаблона — fallback на прямое решение
        resp = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": f"Реши задачу пошагово:\n{task}"}],
            temperature=0
        )
        return {"answer": resp.choices[0].message.content, "template_used": None}

    # Шаг 3: Instantiate — адаптируем шаблон
    steps_text = "\n".join(f"{i+1}. {s}" for i, s in enumerate(template.steps))
    inst_resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": INSTANTIATE_PROMPT.format(
            task_type=template.task_type,
            strategy=template.strategy,
            steps=steps_text,
            task=task
        )}],
        temperature=0
    )
    plan = inst_resp.choices[0].message.content

    # Шаг 4: Reason — решаем по плану
    reason_resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": REASON_PROMPT.format(task=task, plan=plan)}],
        temperature=0
    )
    answer = reason_resp.choices[0].message.content

    return {
        "answer": answer,
        "template_used": template.id,
        "plan": plan,
        "task_type": task_type
    }


# Пример: заполняем буфер и решаем задачу
buffer = MetaBuffer()
buffer.add(ThoughtTemplate(
    id="game24_pairs",
    task_type="arithmetic_puzzle",
    strategy="pair_decomposition",
    steps=[
        "Перечислить все пары чисел и результаты операций (+, -, ×, ÷)",
        "Для каждого промежуточного результата повторить с оставшимися числами",
        "Проверить, даёт ли цепочка целевое число",
        "Вернуть первую найденную цепочку"
    ],
    success_count=15,
    fail_count=2
))

result = buffer_of_thoughts("Числа: 3, 4, 6, 8. Получи 24.", buffer)
print(result["answer"])