Recursion of Thought

from openai import OpenAI
from dataclasses import dataclass, field

client = OpenAI()


@dataclass
class RotFrame:
    """Фрейм стека рекурсии — один уровень задачи."""
    task: str
    depth: int
    result: str | None = None
    subtask_results: dict[str, str] = field(default_factory=dict)


ROT_SYSTEM = """Ты решаешь задачи методом Recursion of Thought.

Если задача сложная — разбей её на подзадачи. Формат:
[GO подзадача_id]: формулировка подзадачи

Если задача простая (решается за 1-2 шага) — реши напрямую.

Когда все подзадачи решены, собери финальный ответ из их результатов.

Правила:
- Каждая подзадача должна быть самодостаточной
- Не включай в подзадачу лишний контекст из родительской задачи
- Подзадач не более 5 за один вызов
- Если получены результаты подзадач — используй их как факты"""


def solve_rot(
    task: str,
    model: str = "gpt-4o",
    max_depth: int = 3,
    depth: int = 0,
) -> str:
    """Рекурсивно решает задачу через RoT."""
    if depth >= max_depth:
        # Базовый случай — решаем напрямую
        resp = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "Реши задачу напрямую, без декомпозиции."},
                {"role": "user", "content": task},
            ],
            temperature=0.3,
        )
        return resp.choices[0].message.content

    # Шаг 1: попросить модель декомпозировать или решить
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": ROT_SYSTEM},
            {"role": "user", "content": task},
        ],
        temperature=0.3,
    )
    answer = resp.choices[0].message.content

    # Шаг 2: извлечь подзадачи [GO ...]
    import re
    subtasks = re.findall(
        r"\[GO\s+(\w+)\]:\s*(.+?)(?=\[GO|\Z)",
        answer,
        re.DOTALL,
    )

    if not subtasks:
        # Модель решила задачу напрямую
        return answer

    # Шаг 3: рекурсивно решить каждую подзадачу
    results = {}
    for sub_id, sub_task in subtasks:
        results[sub_id] = solve_rot(
            sub_task.strip(),
            model=model,
            max_depth=max_depth,
            depth=depth + 1,
        )

    # Шаг 4: собрать результаты и попросить финальный ответ
    return_context = "\n".join(
        f"[RETURN {sid}]: {res}" for sid, res in results.items()
    )

    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": ROT_SYSTEM},
            {
                "role": "user",
                "content": (
                    f"Исходная задача: {task}\n\n"
                    f"Результаты подзадач:\n{return_context}\n\n"
                    "Собери финальный ответ из результатов подзадач."
                ),
            },
        ],
        temperature=0.3,
    )
    return resp.choices[0].message.content


# Пример
result = solve_rot(
    "Сравни TCO (total cost of ownership) за 3 года: "
    "вариант A — собственный GPU-кластер (8×A100, аренда стойки), "
    "вариант B — OpenAI API (500K запросов/день, GPT-4o). "
    "Учти: железо, электричество, персонал, API-стоимость, масштабирование.",
    max_depth=3,
)
print(result)