how to enable function calling in Mistral Agents.py

# -*- coding: utf-8 -*-
"""Custom Function Call.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1--CHjQvfuNC_cWfN-D3ZMzWqKbIRSMLm

In this tutorial, we'll demonstrate how to enable function calling in Mistral Agents using the standard JSON schema format. By defining your function's input parameters with a clear schema, you can make your custom tools seamlessly callable by the agent—enabling powerful, dynamic interactions.

# 1. Setting up dependencies

## 1.1 Installing Mistral library
"""

!pip install mistralai

"""## 1.2 Loading Aviation Stack API
In this section, we’ll integrate the [AviationStack API](https://aviationstack.com/dashboard) to retrieve real-time flight status information. You can sign up for a free API key from their dashboard to get started.
"""

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

"""## 1.3 Loading Mistral API
You can get an API key from https://console.mistral.ai/api-keys
"""

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

"""# 2. Defining the Custom Function

Next, we define a Python function get_flight_status() that calls the AviationStack API to retrieve the real-time status of a flight. The function accepts an optional flight_iata parameter and returns key details such as airline name, flight status, departure and arrival airports, and scheduled times. If no matching flight is found, it gracefully 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."}

get_flight_status('AI101')

"""# 3. Creating the Mistral client and Agent

In this step, we create a Mistral Agent that uses tool-calling to fetch real-time flight information. The agent, named Flight Status Agent, is configured to use the "mistral-medium-2505" model and is equipped with a custom function tool named get_flight_status. This tool is defined using a JSON schema that accepts a single required parameter: the flight's IATA code (e.g., "AI101"). Once deployed, the agent can automatically invoke this function whenever it detects a relevant user query, enabling seamless integration between natural language inputs and structured API responses.
"""

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"]
                }
            }
        }
    ]
)

"""# 4. Starting the Conversation and handling Function Calling

In this step, we initiate a conversation with the Flight Status Agent by asking a natural language question: "What's the current status of AI101?". The Mistral model detects that it should invoke the get_flight_status function and returns a function call request. We parse the arguments, run the function locally using the AviationStack API, and return the result back to the agent using FunctionResultEntry. Finally, the model incorporates the API response and generates a natural language reply with the current flight status, which we print to the console.
"""

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)