MRKL Systems — модульные агенты

import json
import math
from openai import OpenAI

client = OpenAI()


def calculator(expression: str) -> str:
    """Модуль калькулятора — точные вычисления."""
    try:
        # Безопасное вычисление (без eval)
        allowed = {
            "abs": abs, "round": round,
            "min": min, "max": max,
            "sum": sum, "pow": pow,
            "sqrt": math.sqrt, "log": math.log,
        }
        result = eval(expression, {"__builtins__": {}}, allowed)
        return str(result)
    except Exception as e:
        return f"Ошибка вычисления: {e}"


def search(query: str) -> str:
    """Модуль поиска (заглушка — заменить на реальный API)."""
    # В продакшене: Tavily, SerpAPI, Exa
    return f"[Результаты поиска по запросу: '{query}']"


def database_query(query: str) -> str:
    """Модуль базы данных (заглушка)."""
    # В продакшене: SQL-запрос к реальной БД
    return f"[Результат запроса к БД: '{query}']"


MODULES = {
    "calculator": {
        "function": calculator,
        "description": "Точные математические вычисления",
    },
    "search": {
        "function": search,
        "description": "Поиск актуальной информации в интернете",
    },
    "database": {
        "function": database_query,
        "description": "Запрос к базе данных компании",
    },
}


def mrkl_system(
    question: str,
    model: str = "gpt-4o"
) -> dict:
    """
    Базовая MRKL-система: маршрутизация + выполнение + синтез.
    """
    modules_desc = "\n".join(
        f"- {name}: {info['description']}"
        for name, info in MODULES.items()
    )

    # Шаг 1: Маршрутизация — определяем, какие модули нужны
    routing_response = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": (
                f"Доступные модули:\n{modules_desc}\n\n"
                f"Вопрос: {question}\n\n"
                f"Определи, какие модули нужно вызвать и в "
                f"каком порядке. Для каждого вызова укажи "
                f"конкретный запрос.\n\n"
                f"Верни JSON:\n"
                f'{{"steps": ['
                f'{{"module": "имя", "query": "запрос"}}'
                f"]}}"
            )
        }],
        temperature=0,
        response_format={"type": "json_object"},
    )

    try:
        plan = json.loads(
            routing_response.choices[0].message.content
        )
        steps = plan.get("steps", [])
    except (json.JSONDecodeError, KeyError):
        steps = []

    # Шаг 2: Выполнение модулей
    results = []
    for step in steps:
        module_name = step.get("module", "")
        query = step.get("query", "")

        if module_name in MODULES:
            output = MODULES[module_name]["function"](query)
        else:
            output = f"Модуль '{module_name}' не найден"

        results.append({
            "module": module_name,
            "query": query,
            "output": output
        })

    # Шаг 3: Синтез ответа
    results_text = "\n".join(
        f"[{r['module']}] Запрос: {r['query']} → "
        f"Результат: {r['output']}"
        for r in results
    )

    synthesis = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": (
                f"Вопрос: {question}\n\n"
                f"Результаты модулей:\n{results_text}\n\n"
                f"Сформируй финальный ответ на основе "
                f"результатов модулей."
            )
        }],
        temperature=0.5,
    )

    return {
        "question": question,
        "plan": steps,
        "module_results": results,
        "answer": synthesis.choices[0].message.content,
    }


result = mrkl_system(
    "Сколько стоит 15 лицензий Notion по $10/мес на год "
    "в рублях при текущем курсе?"
)
print(result["answer"])

tools = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Выполняет математические вычисления",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": (
                            "Математическое выражение, "
                            "например: '2 * 3 + 4'"
                        ),
                    }
                },
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search",
            "description": "Поиск актуальной информации",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Поисковый запрос",
                    }
                },
                "required": ["query"],
            },
        },
    },
]


def mrkl_with_function_calling(
    question: str,
    model: str = "gpt-4o"
) -> str:
    """MRKL через OpenAI Function Calling (tool use)."""
    messages = [{"role": "user", "content": question}]

    # Цикл вызовов инструментов
    while True:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice="auto",
        )

        message = response.choices[0].message
        messages.append(message)

        # Если модель не вызывает инструменты — готово
        if not message.tool_calls:
            return message.content

        # Выполняем вызовы инструментов
        for tool_call in message.tool_calls:
            fn_name = tool_call.function.name
            fn_args = json.loads(tool_call.function.arguments)

            if fn_name in MODULES:
                fn_result = MODULES[fn_name]["function"](
                    **fn_args
                )
            else:
                fn_result = f"Неизвестный модуль: {fn_name}"

            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": str(fn_result),
            })

import anthropic

client = anthropic.Anthropic()

claude_tools = [
    {
        "name": "calculator",
        "description": "Точные математические вычисления",
        "input_schema": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Математическое выражение",
                }
            },
            "required": ["expression"],
        },
    },
    {
        "name": "search",
        "description": "Поиск актуальной информации",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Поисковый запрос",
                }
            },
            "required": ["query"],
        },
    },
]


def mrkl_claude(question: str) -> str:
    """MRKL через Claude Tool Use."""
    messages = [{"role": "user", "content": question}]

    while True:
        response = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=2048,
            tools=claude_tools,
            messages=messages,
        )

        # Если модель закончила — возвращаем текст
        if response.stop_reason == "end_turn":
            return "".join(
                block.text for block in response.content
                if block.type == "text"
            )

        # Обрабатываем вызовы инструментов
        messages.append({
            "role": "assistant",
            "content": response.content
        })

        tool_results = []
        for block in response.content:
            if block.type == "tool_use":
                if block.name in MODULES:
                    result = MODULES[block.name]["function"](
                        **block.input
                    )
                else:
                    result = f"Модуль {block.name} не найден"

                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": str(result),
                })

        messages.append({"role": "user", "content": tool_results})