agent_orchestration_with_mistral_agents_api.py

# -*- coding: utf-8 -*-
"""Agent Orchestration with Mistral Agents API.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1OV1aXSVYi2l3tRxHo_3e92Q_DBamBdUd

In this tutorial, we'll dive into the powerful Handoffs feature of the Mistral Agents API, which allows you to build dynamic, multi-agent workflows. By enabling agents to delegate tasks to other specialized agents, Handoffs let you create modular, intelligent systems capable of solving complex problems step by step.

# 1. Setting up dependencies

## 1.1 Installing the libraries
"""

!pip install mistralai pydantic

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

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

"""# 2. Agent Prerequisites and Setup

## 2.1 Initializing the Agent
"""

from mistralai import CompletionArgs, ResponseFormat, JSONSchema
from pydantic import BaseModel
from mistralai import Mistral

client = Mistral(MISTRAL_API_KEY)

"""## 2.1 Creating the Custom Function

The adjust_for_inflation function calculates how much a given amount of money would be worth after accounting for inflation over time. It uses the compound formula based on the number of years and the annual inflation rate. If the end year is before the start year, it returns an error. Otherwise, it returns the adjusted value along with the input details. For example, adjust_for_inflation(1000, 1899, 2025, 10) shows what ₹1000 from 1899 would be worth in 2025 at 10% inflation.
"""

def adjust_for_inflation(amount: float, start_year: int, end_year: int, annual_inflation_rate: float):
    """
    Calculates inflation-adjusted value using compound formula.
    """
    if end_year < start_year:
        return {"error": "End year must be greater than or equal to start year."}

    years = end_year - start_year
    adjusted_value = amount * ((1 + annual_inflation_rate / 100) ** years)

    return {
        "original_amount": amount,
        "start_year": start_year,
        "end_year": end_year,
        "inflation_rate": annual_inflation_rate,
        "adjusted_value": round(adjusted_value, 2)
    }

adjust_for_inflation(1000, 1899, 2025, 10)

"""## 2.2 Creating Structured Output for Mathematical Reasoning

We're using the inflation_tool definition to enable Mistral Agents to call a custom function (adjust_for_inflation) that calculates how the value of money changes over time due to inflation. The schema clearly defines the required inputs—amount, years, and inflation rate—so the agent can interpret user queries and trigger this function when needed. This setup allows seamless integration of real-world calculations into conversations.
"""

class CalcResult(BaseModel):
    reasoning: str
    result: str

inflation_tool = {
    "type": "function",
    "function": {
        "name": "adjust_for_inflation",
        "description": "Calculate the value of money adjusted for inflation over a time period.",
        "parameters": {
            "type": "object",
            "properties": {
                "amount": {
                    "type": "number",
                    "description": "Original amount of money"
                },
                "start_year": {
                    "type": "integer",
                    "description": "The starting year for inflation adjustment"
                },
                "end_year": {
                    "type": "integer",
                    "description": "The ending year for inflation adjustment"
                },
                "annual_inflation_rate": {
                    "type": "number",
                    "description": "Annual inflation rate in percent"
                }
            },
            "required": ["amount", "start_year", "end_year", "annual_inflation_rate"]
        }
    }
}

"""# 3. Creating the Agents

## 3.1 Defining the different agents

In this setup, we define a multi-agent system using Mistral Agents API to handle inflation-related economic queries. The main agent (economics-agent) acts as a coordinator that routes tasks to specialized agents. The inflation-agent performs inflation adjustment calculations using a custom function. If the inflation rate is missing from the query, the websearch-agent fetches it from the internet. The calculator-agent handles complex numerical computations with step-by-step reasoning, while the graph-agent uses the code interpreter to visualize inflation trends over time. Together, these agents collaborate via handoffs to deliver accurate, dynamic responses to economic queries.
"""

# Main Agent
economics_agent = client.beta.agents.create(
    model="mistral-large-latest",
    name="economics-agent",
    description="Handles economic queries and delegates inflation calculations.",
)

# Inflation Function Agent
inflation_agent = client.beta.agents.create(
    model="mistral-large-latest",
    name="inflation-agent",
    description="Agent that calculates inflation-adjusted value using a custom function.",
    tools=[inflation_tool],
)

# Web Search Agent
websearch_agent = client.beta.agents.create(
    model="mistral-large-latest",
    name="websearch-agent",
    description="Agent that can search the internet for missing economic data such as inflation rates.",
    tools=[{"type": "web_search"}]
)


# Calculator Agent
from pydantic import BaseModel

class CalcResult(BaseModel):
    reasoning: str
    result: str

calculator_agent = client.beta.agents.create(
    model="mistral-large-latest",
    name="calculator-agent",
    description="Agent used to make detailed calculations.",
    instructions="When doing calculations, explain step by step.",
    completion_args=CompletionArgs(
        response_format=ResponseFormat(
            type="json_schema",
            json_schema=JSONSchema(
                name="calc_result",
                schema=CalcResult.model_json_schema(),
            )
        )
    )
)

# Graph Agent
graph_agent = client.beta.agents.create(
    model="mistral-large-latest",
    name="graph-agent",
    description="Agent that generates graphs using code interpreter.",
    instructions="Use code interpreter to draw inflation trends.",
    tools=[{"type": "code_interpreter"}]
)

"""## 3.2 Defining the Handoffs Responsibilities

This configuration defines how agents delegate tasks among each other:

* The Main Agent (economics_agent) serves as the entry point and delegates queries either to the inflation_agent (for inflation calculations) or the websearch_agent (to fetch missing data like inflation rates).

* The inflation_agent, after receiving either the user query or web-fetched data, can further pass tasks to the calculator_agent (for detailed math) or graph_agent (to visualize trends).

* The websearch_agent can pass control to the inflation_agent after retrieving required information, like the inflation rate.

* calculator_agent and graph_agent are considered terminal agents. However, optional mutual handoff is enabled in case one needs to do follow-up work (e.g., graphing a calculated result or vice versa).
"""

# Main Agent hands off to inflation_agent and websearch_agent
economics_agent = client.beta.agents.update(
    agent_id=economics_agent.id,
    handoffs=[inflation_agent.id, websearch_agent.id]
)

# Inflation Agent can delegate to calculator_agent or graph_agent if deeper analysis or visualization is needed
inflation_agent = client.beta.agents.update(
    agent_id=inflation_agent.id,
    handoffs=[calculator_agent.id, graph_agent.id]
)

# Web Search Agent can hand off to inflation_agent (after finding the missing rate)
websearch_agent = client.beta.agents.update(
    agent_id=websearch_agent.id,
    handoffs=[inflation_agent.id]
)

# Calculator and Graph agents are terminal—they don't hand off further
# But if needed, we could let them hand off to each other:
calculator_agent = client.beta.agents.update(
    agent_id=calculator_agent.id,
    handoffs=[graph_agent.id]  # Optional
)

graph_agent = client.beta.agents.update(
    agent_id=graph_agent.id,
    handoffs=[calculator_agent.id]  # Optional
)

"""# 4. Running the Agent

## 4.1 Example A: What is the current inflation rate in India?

In this example, the prompt "What is the current inflation rate in India?" is passed to the economics_agent, which is the main entry point for handling economic queries. Since the question requires real-time data that isn't included in the agent’s static knowledge, the economics_agent automatically hands off the query to the websearch_agent, which is equipped with web search capabilities.
"""

prompt = "What is the current inflation rate in India?"
response = client.beta.conversations.start(
    agent_id=economics_agent.id,
    inputs=prompt
)

response

print(response.outputs[-1].content[0].text)

"""## 4.2 Example B: What is the inflation-adjusted value of 5,000 from the year 2010 to 2023 with annual inflation rate of 6.5%. Explain calculation steps and plot a graph with data labels

This code block sends the prompt to an economics agent, checks if the agent triggers a specific function call (adjust_for_inflation), executes that function locally with the provided arguments, and then returns the computed result back to the agent. Finally, it prints the agent's response, which includes the inflation calculation explanation, along with the Python code to plot the trend.
"""

import json
from mistralai.models import FunctionResultEntry

prompt = """What is the inflation-adjusted value of 5,000 from the year 2010 to 2023 with annual inflation rate of 6.5%.
Explain calculation steps and plot a graph with data labels"""

response = client.beta.conversations.start(
    agent_id=economics_agent.id,
    inputs=prompt
)

# Check for function call
if response.outputs[-1].type == "function.call" and response.outputs[-1].name == "adjust_for_inflation":
    args = json.loads(response.outputs[-1].arguments)

    # Run local function
    function_result = json.dumps(adjust_for_inflation(**args))

    # Return result to Mistral
    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)

"""Execcuting the following code plots the inflation-adjusted value of $5,000 from 2010 to 2023"""

import matplotlib.pyplot as plt
import numpy as np

# Parameters
original_amount = 5000
start_year = 2010
end_year = 2023
inflation_rate = 6.5 / 100  # Convert percentage to decimal

# Calculate the number of years
num_years = end_year - start_year + 1

# Calculate the adjusted value for each year
years = np.arange(start_year, end_year + 1)
adjusted_values = original_amount * (1 + inflation_rate) ** (years - start_year)

# Plot the graph
plt.figure(figsize=(10, 6))
plt.plot(years, adjusted_values, marker='o', linestyle='-', color='b')

# Add data labels
for year, value in zip(years, adjusted_values):
    plt.text(year, value, f'${value:.2f}', ha='right')

# Add titles and labels
plt.title('Inflation-Adjusted Value Over Time')
plt.xlabel('Year')
plt.ylabel('Adjusted Value')

# Save the plot as an image
plt.savefig('inflation_adjusted_value.png')

# Show the plot
plt.show()