Включение динамического вызова функций в агентах Mistral через JSON Schema

Использование JSON Schema для вызова функций в агентах Mistral

В этом руководстве показано, как включить вызов функций в агентах Mistral, определяя входные параметры функций с помощью стандартного формата JSON Schema. Этот метод позволяет вашим пользовательским инструментам быть легко вызываемыми агентом, обеспечивая мощное и динамическое взаимодействие.

Интеграция API AviationStack для получения статуса рейса

Пример демонстрирует интеграцию с API AviationStack для получения данных о статусе рейса в реальном времени. Это показывает, как внешние API можно использовать как вызываемые функции внутри агента Mistral.

Шаг 1: Установка зависимостей

Сначала установите библиотеку Mistral:

pip install mistralai

Безопасно загрузите ключ API Mistral:

from getpass import getpass
MISTRAL_API_KEY = getpass('Введите ключ API Mistral: ')

Аналогично получите и загрузите ключ API AviationStack:

AVIATIONSTACK_API_KEY = getpass('Введите ключ API AviationStack: ')

Шаг 2: Определение пользовательской функции

Определите функцию Python get_flight_status(), которая обращается к API AviationStack. Она принимает необязательный параметр flight_iata и возвращает такие данные, как название авиакомпании, статус рейса, аэропорты вылета и прилета, а также расписание. Если рейс не найден, возвращается сообщение об ошибке.

import requests
from typing import Dict
def get_flight_status(flight_iata=None):
    """
    Получить статус рейса с использованием необязательных фильтров: dep_iata, arr_iata, flight_iata.
    """
    params = {
        "access_key": AVIATIONSTACK_API_KEY,
        "flight_iata": flight_iata  
    }
 
    response = requests.get("http://api.aviationstack.com/v1/flights", params=params)
    data = response.json()
 
    if "data" in data and data["data"]:
        flight = data["data"][0]
        return {
            "airline": flight["airline"]["name"],
            "flight_iata": flight["flight"]["iata"],
            "status": flight["flight_status"],
            "departure_airport": flight["departure"]["airport"],
            "arrival_airport": flight["arrival"]["airport"],
            "scheduled_departure": flight["departure"]["scheduled"],
            "scheduled_arrival": flight["arrival"]["scheduled"],
        }
    else:
        return {"error": "Рейс с указанными параметрами не найден."}

Шаг 3: Создание клиента Mistral и агента

Создайте агента Mistral, оснащённого пользовательской функцией, определённой через JSON Schema, которая принимает код IATA рейса. Агент использует модель "mistral-medium-2505" и интегрирует функцию get_flight_status.

from mistralai import Mistral
client = Mistral(MISTRAL_API_KEY)
 
flight_status_agent = client.beta.agents.create(
    model="mistral-medium-2505",
    description="Предоставляет статус рейса в реальном времени с использованием API aviationstack.",
    name="Flight Status Agent",
    tools=[
        {
            "type": "function",
            "function": {
                "name": "get_flight_status",
                "description": "Получить текущий статус рейса по коду IATA (например, AI101).",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "flight_iata": {
                            "type": "string",
                            "description": "Код IATA рейса (например, AI101)"
                        },
                    },
                    "required": ["flight_iata"]
                }
            }
        }
    ]
)

Шаг 4: Взаимодействие и обработка вызова функции

Начните разговор с агентом, задав вопрос о статусе рейса. Агент обнаруживает необходимость вызова функции get_flight_status, выполняет её и возвращает отформатированный ответ о статусе рейса.

from mistralai import FunctionResultEntry
import json
 
# Пользователь начинает разговор
response = client.beta.conversations.start(
    agent_id=flight_status_agent.id,
    inputs=[{"role": "user", "content": "Каков текущий статус рейса AI101?"}]
)
 
# Проверяем, запросил ли модель вызов функции
if response.outputs[-1].type == "function.call" and response.outputs[-1].name == "get_flight_status":
    args = json.loads(response.outputs[-1].arguments)
 
    # Выполняем функцию
    function_result = json.dumps(get_flight_status(**args))
 
    # Создаём запись с результатом
    result_entry = FunctionResultEntry(
        tool_call_id=response.outputs[-1].tool_call_id,
        result=function_result
    )
 
    # Отправляем результат обратно агенту
    response = client.beta.conversations.append(
        conversation_id=response.conversation_id,
        inputs=[result_entry]
    )
 
    print(response.outputs[-1].content)
else:
    print(response.outputs[-1].content)

Этот пример показывает, как естественные языковые запросы могут триггерить структурированные вызовы API через агентов Mistral, обеспечивая динамичное получение данных в реальном времени.

Полный пример и ноутбук доступны в репозитории на GitHub по ссылке из оригинальной статьи.