Enabling Dynamic Function Calling in Mistral Agents via JSON Schema

Using JSON Schema for Function Calling in Mistral Agents

This tutorial illustrates how to enable function calling in Mistral Agents by defining your function's input parameters using the standard JSON schema format. Leveraging this method allows your custom tools to be seamlessly callable by the agent, facilitating powerful and dynamic interactions.

Integrating AviationStack API to Retrieve Flight Status

The example demonstrates integration with the AviationStack API to fetch real-time flight status data. This showcases how external APIs can be incorporated as callable functions within a Mistral Agent.

Step 1: Setting up Dependencies

First, install the Mistral library:

pip install mistralai

Load your Mistral API key securely:

from getpass import getpass
MISTRAL_API_KEY = getpass('Enter Mistral API Key: ')

Similarly, obtain and load your AviationStack API key:

AVIATIONSTACK_API_KEY = getpass('Enter Aviation Stack API: ')

Step 2: Defining the Custom Function

Define a Python function get_flight_status() which queries the AviationStack API. It accepts an optional flight_iata parameter and returns details such as airline name, flight status, departure and arrival airports, and scheduled times. If no flight matches, it returns an error message.

import requests
from typing import Dict
def get_flight_status(flight_iata=None):
    """
    Retrieve flight status using optional filters: 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": "No flight found for the provided parameters."}

Step 3: Creating the Mistral Client and Agent

Create a Mistral Agent equipped with a custom function tool that uses JSON schema to accept the flight's IATA code. The agent uses the "mistral-medium-2505" model and integrates the get_flight_status function.

from mistralai import Mistral
client = Mistral(MISTRAL_API_KEY)
 
flight_status_agent = client.beta.agents.create(
    model="mistral-medium-2505",
    description="Provides real-time flight status using aviationstack API.",
    name="Flight Status Agent",
    tools=[
        {
            "type": "function",
            "function": {
                "name": "get_flight_status",
                "description": "Retrieve the current status of a flight by its IATA code (e.g. AI101).",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "flight_iata": {
                            "type": "string",
                            "description": "IATA code of the flight (e.g. AI101)"
                        },
                    },
                    "required": ["flight_iata"]
                }
            }
        }
    ]
)

Step 4: Interaction and Function Calling Handling

Start a conversation with the agent by asking about a flight status. The agent detects the need to call the get_flight_status function, executes it, and returns the formatted flight status.

from mistralai import FunctionResultEntry
import json
 
# User starts a conversation
response = client.beta.conversations.start(
    agent_id=flight_status_agent.id,
    inputs=[{"role": "user", "content": "What's the current status of AI101?"}]
)
 
# Check if model requested a function call
if response.outputs[-1].type == "function.call" and response.outputs[-1].name == "get_flight_status":
    args = json.loads(response.outputs[-1].arguments)
 
    # Run the function
    function_result = json.dumps(get_flight_status(**args))
 
    # Create result entry
    result_entry = FunctionResultEntry(
        tool_call_id=response.outputs[-1].tool_call_id,
        result=function_result
    )
 
    # Return result to agent
    response = client.beta.conversations.append(
        conversation_id=response.conversation_id,
        inputs=[result_entry]
    )
 
    print(response.outputs[-1].content)
else:
    print(response.outputs[-1].content)

This approach highlights how natural language queries can trigger structured API calls through Mistral Agents, enabling dynamic and real-time data retrieval.

For the complete example and notebook, check the GitHub repository linked in the original post.