Thought Propagation

from openai import OpenAI

client = OpenAI()


def thought_propagation(
    task: str,
    n_analogies: int = 3,
    title: str = "gpt-4o",
) -> dict:
    """
    Thought Propagation: decompose → solve analogies → propagate.
    """
    # Шаг 1: Найти аналогичные подзадачи
    analogies_response = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "system",
                "content": (
                    "Ты — эксперт по декомпозиции задач. "
                    "Найди аналогичные подзадачи, которые "
                    "проще исходной, но имеют ту же структуру."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Задача: {task}\n\n"
                    f"Предложи {n_analogies} аналогичных подзадач, "
                    "которые проще исходной, но имеют ту же "
                    "структуру решения. Для каждой объясни, "
                    "почему она аналогична."
                ),
            },
        ],
        temperature=0.3,
    )
    analogies = analogies_response.choices[0].message.content

    # Шаг 2: Решить аналогичные подзадачи
    solutions_response = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "system",
                "content": (
                    "Ты — эксперт по решению задач. "
                    "Реши каждую подзадачу пошагово."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Исходная задача: {task}\n\n"
                    f"Аналогичные подзадачи:\n{analogies}\n\n"
                    "Реши каждую подзадачу пошагово. "
                    "Выяви паттерн решения."
                ),
            },
        ],
        temperature=0.2,
    )
    solutions = solutions_response.choices[0].message.content

    # Шаг 3: Распространить решения
    propagation_response = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "system",
                "content": (
                    "Ты — эксперт по синтезу решений. "
                    "Перенеси решения подзадач на исходную "
                    "задачу, адаптируя и комбинируя их."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Исходная задача: {task}\n\n"
                    f"Решения аналогичных подзадач:\n{solutions}\n\n"
                    "Перенеси (propagate) решения подзадач на "
                    "исходную задачу. Объясни, как каждое "
                    "решение подзадачи вносит вклад в общий ответ."
                ),
            },
        ],
        temperature=0.2,
    )

    return {
        "analogies": analogies,
        "solutions": solutions,
        "final_answer": propagation_response.choices[0].message.content,
        "total_tokens": sum(
            r.usage.total_tokens
            for r in [
                analogies_response,
                solutions_response,
                propagation_response,
            ]
        ),
    }


result = thought_propagation(
    task=(
        "Найди кратчайший путь от узла S до узла T "
        "в графе: S-A:3, S-B:5, A-C:2, A-D:4, "
        "B-D:1, B-E:3, C-T:6, D-T:2, E-T:5."
    ),
)
print(result["final_answer"])

def iterative_thought_propagation(
    task: str,
    max_iterations: int = 2,
    title: str = "gpt-4o",
) -> dict:
    """
    Итеративный TP: после первого прохода модель
    оценивает качество решения и при необходимости
    генерирует дополнительные аналогии.
    """
    client = OpenAI()
    history = []

    for i in range(max_iterations):
        context = (
            f"\n\nПредыдущие попытки:\n"
            + "\n".join(history)
            if history
            else ""
        )

        response = client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": (
                        f"Задача: {task}{context}\n\n"
                        "1. Предложи 2-3 аналогичные подзадачи "
                        "(проще, но с той же структурой).\n"
                        "2. Реши каждую.\n"
                        "3. Перенеси решения на исходную задачу.\n"
                        "4. Оцени уверенность (0-100%).\n\n"
                        "Если уверенность < 80%, укажи, "
                        "какие аспекты требуют дополнительных "
                        "аналогий."
                    ),
                }
            ],
            temperature=0.3,
        )

        answer = response.choices[0].message.content
        history.append(f"Итерация {i + 1}: {answer[:500]}...")

        if "уверенность" in answer.lower() and any(
            f"{n}%" in answer for n in range(80, 101)
        ):
            break

    return {"answer": answer, "iterations": len(history)}