Self-Ask: декомпозиция вопроса на подвопросы

from dataclasses import dataclass, field
import anthropic

client = anthropic.Anthropic()


@dataclass
class SubQuestion:
    id: int
    question: str
    depends_on: list[int]
    answer: str = ""


DECOMPOSE_TOOL = {
    "name": "decompose_question",
    "description": (
        "Декомпозируй сложный вопрос на простые атомарные подвопросы. "
        "Каждый подвопрос должен отвечаться за один поиск или вычисление. "
        "Если вопрос простой и ответ известен напрямую — верни is_complex=false."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "is_complex": {
                "type": "boolean",
                "description": "True если вопрос требует нескольких промежуточных фактов"
            },
            "sub_questions": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "id": {"type": "integer"},
                        "question": {
                            "type": "string",
                            "description": "Атомарный вопрос с одним конкретным ответом"
                        },
                        "depends_on": {
                            "type": "array",
                            "items": {"type": "integer"},
                            "description": "ID подвопросов, чьи ответы нужны для этого"
                        }
                    },
                    "required": ["id", "question", "depends_on"]
                }
            }
        },
        "required": ["is_complex", "sub_questions"]
    }
}


def decompose(question: str) -> list[SubQuestion]:
    """Возвращает список подвопросов или пустой список если вопрос простой."""
    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=512,
        system=(
            "Ты аналитик, декомпозирующий multi-hop вопросы.\n"
            "Подвопрос атомарный = один факт, одна дата, одно имя.\n"
            "Каждый depends_on точно отражает реальную зависимость: "
            "подвопрос 3 зависит от 1 только если ответ 1 входит в текст вопроса 3."
        ),
        tools=[DECOMPOSE_TOOL],
        tool_choice={"type": "tool", "name": "decompose_question"},
        messages=[{"role": "user", "content": question}]
    )
    for block in response.content:
        if block.type == "tool_use" and block.name == "decompose_question":
            data = block.input
            if not data.get("is_complex"):
                return []
            return [
                SubQuestion(
                    id=sq["id"],
                    question=sq["question"],
                    depends_on=sq.get("depends_on", [])
                )
                for sq in data.get("sub_questions", [])
            ]
    return []

def answer_sub_question(sub_q: SubQuestion, answers: dict[int, str]) -> str:
    """Отвечает на один подвопрос, передавая промежуточные ответы как контекст."""
    context = ""
    if sub_q.depends_on:
        deps = "\n".join(
            f"  - Подвопрос {dep_id}: {answers[dep_id]}"
            for dep_id in sub_q.depends_on
            if dep_id in answers
        )
        context = f"\n\nПромежуточные ответы:\n{deps}"

    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=200,
        system="Отвечай кратко и конкретно: только факт, без лишних слов.",
        messages=[{
            "role": "user",
            "content": f"Вопрос: {sub_q.question}{context}"
        }]
    )
    return response.content[0].text.strip()


def synthesize(question: str, sub_qs: list[SubQuestion]) -> str:
    """Собирает финальный ответ из промежуточных ответов."""
    facts = "\n".join(
        f"  {sq.id}. {sq.question}\n     Ответ: {sq.answer}"
        for sq in sub_qs
    )
    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=1024,
        system="Синтезируй собранные факты в связный ответ на исходный вопрос.",
        messages=[{
            "role": "user",
            "content": f"Исходный вопрос: {question}\n\nСобранные факты:\n{facts}"
        }]
    )
    return response.content[0].text


def self_ask(question: str) -> str:
    """
    Self-Ask агент:
    1. Декомпозирует вопрос на атомарные подвопросы
    2. Отвечает на каждый по порядку (с учётом зависимостей)
    3. Синтезирует финальный ответ
    """
    # Шаг 1: Декомпозиция
    sub_qs = decompose(question)

    if not sub_qs:
        # Простой вопрос — отвечаем напрямую
        simple = SubQuestion(0, question, [])
        return answer_sub_question(simple, {})

    print(f"[Self-Ask] {len(sub_qs)} подвопросов:")

    # Шаг 2: Ответы в порядке зависимостей
    answers: dict[int, str] = {}
    for sq in topological_sort(sub_qs):
        print(f"  [{sq.id}] {sq.question}")
        sq.answer = answer_sub_question(sq, answers)
        answers[sq.id] = sq.answer
        print(f"       → {sq.answer[:80]}")

    # Шаг 3: Синтез
    return synthesize(question, sub_qs)


# Пример использования
result = self_ask(
    "Кто дольше прожил — автор «Мастера и Маргариты» или «Войны и мира»?"
)
print(result)

SEARCH_TOOL = {
    "name": "web_search",
    "description": "Поиск актуальной информации в интернете",
    "input_schema": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Краткий поисковый запрос, 3–7 слов"
            }
        },
        "required": ["query"]
    }
}


def answer_with_search(sub_q: SubQuestion, answers: dict[int, str]) -> str:
    """Отвечает на подвопрос с возможностью поиска (один шаг tool use)."""
    context = ""
    if sub_q.depends_on:
        deps = "\n".join(
            f"  - Подвопрос {d}: {answers[d]}"
            for d in sub_q.depends_on if d in answers
        )
        context = f"\n\nКонтекст:\n{deps}"

    messages = [{
        "role": "user",
        "content": f"{sub_q.question}{context}"
    }]

    # Первый вызов: модель решает — использовать поиск или ответить напрямую
    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=256,
        system=(
            "Отвечай кратко. Для точных фактов, дат и актуальных данных "
            "используй поиск. Для простых вычислений — отвечай напрямую."
        ),
        tools=[SEARCH_TOOL],
        messages=messages
    )

    # Если модель вызвала поиск — выполняем и получаем финальный ответ
    for block in response.content:
        if block.type == "tool_use" and block.name == "web_search":
            search_result = run_search(block.input["query"])  # ваша реализация поиска
            messages.append({"role": "assistant", "content": response.content})
            messages.append({"role": "user", "content": [
                {"type": "tool_result", "tool_use_id": block.id, "content": search_result}
            ]})
            final = client.messages.create(
                model="claude-opus-4-6",
                max_tokens=200,
                tools=[SEARCH_TOOL],
                messages=messages
            )
            return final.content[0].text.strip()

    # Модель ответила без поиска
    return response.content[0].text.strip()


def run_search(query: str) -> str:
    """Заглушка — замените реальным поиском (DuckDuckGo, Tavily, SerpAPI и т.д.)."""
    # import httpx
    # resp = httpx.get("https://api.tavily.com/search", params={"query": query, "api_key": ...})
    # return resp.json()["results"][0]["content"]
    return f"[результаты поиска для: {query}]"

def topological_sort(sub_qs: list[SubQuestion]) -> list[SubQuestion]:
    """Сортирует подвопросы в порядке зависимостей (зависимости — раньше)."""
    id_to_sq = {sq.id: sq for sq in sub_qs}
    visited: set[int] = set()
    result: list[SubQuestion] = []

    def visit(sq_id: int) -> None:
        if sq_id in visited:
            return
        visited.add(sq_id)
        for dep_id in id_to_sq[sq_id].depends_on:
            if dep_id in id_to_sq:
                visit(dep_id)
        result.append(id_to_sq[sq_id])

    for sq in sub_qs:
        visit(sq.id)

    return result

import asyncio


async def answer_async(sub_q: SubQuestion, answers: dict[int, str]) -> tuple[int, str]:
    """Async-обёртка над answer_sub_question для параллельного выполнения."""
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(None, answer_sub_question, sub_q, answers)
    return sub_q.id, result


async def self_ask_parallel(question: str) -> str:
    """Self-Ask с параллельным выполнением независимых подвопросов."""
    sub_qs = decompose(question)
    if not sub_qs:
        return answer_sub_question(SubQuestion(0, question, []), {})

    answers: dict[int, str] = {}
    completed: set[int] = set()
    id_to_sq = {sq.id: sq for sq in sub_qs}

    while len(completed) < len(sub_qs):
        # Находим подвопросы, все зависимости которых уже отвечены
        ready = [
            sq for sq in sub_qs
            if sq.id not in completed
            and all(dep in completed for dep in sq.depends_on)
        ]

        if not ready:
            break  # цикл зависимостей — не должно случаться с корректным DAG

        # Отвечаем на все готовые подвопросы параллельно
        tasks = [answer_async(sq, answers) for sq in ready]
        results = await asyncio.gather(*tasks)

        for sq_id, answer in results:
            answers[sq_id] = answer
            id_to_sq[sq_id].answer = answer
            completed.add(sq_id)

    return synthesize(question, sub_qs)


# Запуск
result = asyncio.run(self_ask_parallel(
    "Кто дольше прожил — автор «Мастера и Маргариты» или «Войны и мира»?"
))