Mixture-of-Agents (MoA)

import asyncio
from dataclasses import dataclass
from openai import AsyncOpenAI

@dataclass
class MoAConfig:
    """Конфигурация Mixture-of-Agents."""
    proposers: list[str]       # модели для генерации
    aggregator: str            # модель для синтеза
    n_layers: int = 2          # количество слоёв (1 = proposers + aggregator)
    temperature: float = 0.7


async def call_model(
    client: AsyncOpenAI,
    model: str,
    system: str,
    user: str,
    temperature: float = 0.7
) -> str:
    """Вызов одной модели."""
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user}
        ],
        temperature=temperature
    )
    return resp.choices[0].message.content


PROPOSER_SYSTEM = "Дай подробный, вдумчивый ответ на вопрос."

AGGREGATOR_SYSTEM = """Ты — aggregator в системе Mixture-of-Agents.
Тебе даны ответы нескольких моделей на один вопрос.
Синтезируй лучший финальный ответ:
- Возьми сильные стороны каждого ответа
- Исправь ошибки и неточности
- Дай связный, структурированный итог
- Не упоминай, что ты синтезируешь ответы других моделей"""

REFINER_SYSTEM = """Тебе даны несколько ответов на вопрос.
Улучши их: исправь неточности, добавь пропущенное, убери избыточное.
Дай свой улучшенный ответ."""


async def moa_layer(
    client: AsyncOpenAI,
    task: str,
    models: list[str],
    prev_responses: list[str] | None,
    system: str,
    temperature: float
) -> list[str]:
    """Один слой MoA: параллельные вызовы моделей."""
    if prev_responses:
        refs = "\n\n".join(
            f"--- Ответ {i+1} ---\n{r}"
            for i, r in enumerate(prev_responses)
        )
        user_prompt = f"Вопрос: {task}\n\nОтветы предыдущего слоя:\n{refs}\n\nДай свой улучшенный ответ."
    else:
        user_prompt = task

    tasks = [
        call_model(client, model, system, user_prompt, temperature)
        for model in models
    ]
    return await asyncio.gather(*tasks)


async def mixture_of_agents(
    task: str,
    config: MoAConfig,
    base_url: str | None = None
) -> str:
    """Полный pipeline Mixture-of-Agents."""
    client = AsyncOpenAI(base_url=base_url)

    responses = None

    # Слои proposers / refiners
    for layer in range(config.n_layers):
        system = PROPOSER_SYSTEM if layer == 0 else REFINER_SYSTEM
        responses = await moa_layer(
            client, task, config.proposers,
            responses, system, config.temperature
        )

    # Финальная агрегация
    refs = "\n\n".join(
        f"--- Ответ {i+1} ---\n{r}"
        for i, r in enumerate(responses)
    )
    final = await call_model(
        client, config.aggregator, AGGREGATOR_SYSTEM,
        f"Вопрос: {task}\n\nОтветы:\n{refs}",
        temperature=0.3
    )

    return final


# Пример использования
async def main():
    config = MoAConfig(
        proposers=[
            "meta-llama/Llama-3-70b-chat-hf",
            "Qwen/Qwen2.5-72B-Instruct",
            "mistralai/Mixtral-8x22B-Instruct-v0.1"
        ],
        aggregator="meta-llama/Llama-3-70b-chat-hf",
        n_layers=2,
        temperature=0.7
    )

    result = await mixture_of_agents(
        task="Объясни разницу между RAG и fine-tuning для продакшена",
        config=config,
        base_url="https://api.together.xyz/v1"  # Together AI API
    )
    print(result)

# asyncio.run(main())

async def budget_moa(
    task: str,
    model: str = "gpt-4o-mini",
    n_proposers: int = 5,
    temperatures: list[float] | None = None
) -> str:
    """
    Бюджетный MoA: одна модель с разными temperature.
    Дешевле полноценного MoA, но даёт часть выигрыша.
    """
    client = AsyncOpenAI()
    temps = temperatures or [0.3, 0.5, 0.7, 0.9, 1.1]

    # Генерация разнообразных ответов
    tasks = [
        call_model(client, model, PROPOSER_SYSTEM, task, t)
        for t in temps[:n_proposers]
    ]
    proposals = await asyncio.gather(*tasks)

    # Агрегация
    refs = "\n\n".join(
        f"--- Вариант {i+1} ---\n{r}"
        for i, r in enumerate(proposals)
    )
    return await call_model(
        client, model, AGGREGATOR_SYSTEM,
        f"Вопрос: {task}\n\nВарианты:\n{refs}",
        temperature=0.2
    )