Function Calling Mistral Agent¶
This notebook shows you how to use our Mistral agent, powered by function calling capabilities.
Initial Setup¶
Let's start by importing some simple building blocks.
The main thing we need is:
- the OpenAI API (using our own
llama_index
LLM class) - a place to keep conversation history
- a definition for tools that our agent can use.
If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.
%pip install llama-index-llms-mistralai
%pip install llama-index-embeddings-mistralai
!pip install llama-index
import json
from typing import Sequence, List
from llama_index.llms.mistralai import MistralAI
from llama_index.core.llms import ChatMessage
from llama_index.core.tools import BaseTool, FunctionTool
import nest_asyncio
nest_asyncio.apply()
Let's define some very simple calculator tools for our agent.
def multiply(a: int, b: int) -> int:
"""Multiple two integers and returns the result integer"""
return a * b
multiply_tool = FunctionTool.from_defaults(fn=multiply)
def add(a: int, b: int) -> int:
"""Add two integers and returns the result integer"""
return a + b
add_tool = FunctionTool.from_defaults(fn=add)
Make sure your MISTRAL_API_KEY is set. Otherwise explicitly specify the api_key
parameter.
llm = MistralAI(model="mistral-large-latest")
Initialize Mistral Agent¶
Here we initialize a simple Mistral agent with calculator functions.
from llama_index.core.agent import FunctionCallingAgent
agent = FunctionCallingAgent.from_tools(
[multiply_tool, add_tool],
llm=llm,
verbose=True,
allow_parallel_tool_calls=False,
)
Chat¶
response = agent.chat("What is (121 + 2) * 5?")
print(str(response))
Added user message to memory: What is (121 + 2) * 5? === Calling Function === Calling function: add with args: {"a": 121, "b": 2} === Calling Function === Calling function: multiply with args: {"a": 123, "b": 5} assistant: The result of (121 + 2) * 5 is 615.
# inspect sources
print(response.sources)
[ToolOutput(content='123', tool_name='add', raw_input={'args': (), 'kwargs': {'a': 121, 'b': 2}}, raw_output=123), ToolOutput(content='615', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 123, 'b': 5}}, raw_output=615)]
Async Chat¶
Also let's re-enable parallel function calling so that we can call two multiply
operations simultaneously.
# enable parallel function calling
agent = FunctionCallingAgent.from_tools(
[multiply_tool, add_tool],
llm=llm,
verbose=True,
allow_parallel_tool_calls=True,
)
response = await agent.achat("What is (121 * 3) + (5 * 8)?")
print(str(response))
Added user message to memory: What is (121 * 3) + (5 * 8)? === Calling Function === Calling function: multiply with args: {"a": 121, "b": 3} === Calling Function === Calling function: multiply with args: {"a": 5, "b": 8} === Calling Function === Calling function: add with args: {"a": 363, "b": 40} assistant: The result of (121 * 3) + (5 * 8) is 403.
Mistral Agent over RAG Pipeline¶
Build a Mistral agent over a simple 10K document. We use both Mistral embeddings and mistral-medium to construct the RAG pipeline, and pass it to the Mistral agent as a tool.
!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'
--2024-03-23 11:13:41-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8002::154, 2606:50c0:8001::154, ... Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected. HTTP request sent, awaiting response... 200 OK Length: 1880483 (1.8M) [application/octet-stream] Saving to: ‘data/10k/uber_2021.pdf’ data/10k/uber_2021. 100%[===================>] 1.79M --.-KB/s in 0.09s 2024-03-23 11:13:41 (19.3 MB/s) - ‘data/10k/uber_2021.pdf’ saved [1880483/1880483]
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.llms.mistralai import MistralAI
embed_model = MistralAIEmbedding()
query_llm = MistralAI(model="mistral-medium")
# load data
uber_docs = SimpleDirectoryReader(
input_files=["./data/10k/uber_2021.pdf"]
).load_data()
# build index
uber_index = VectorStoreIndex.from_documents(
uber_docs, embed_model=embed_model
)
uber_engine = uber_index.as_query_engine(similarity_top_k=3, llm=query_llm)
query_engine_tool = QueryEngineTool(
query_engine=uber_engine,
metadata=ToolMetadata(
name="uber_10k",
description=(
"Provides information about Uber financials for year 2021. "
"Use a detailed plain text question as input to the tool."
),
),
)
from llama_index.core.agent import FunctionCallingAgent
agent = FunctionCallingAgent.from_tools(
[query_engine_tool], llm=llm, verbose=True
)
response = agent.chat(
"Tell me both the risk factors and tailwinds for Uber? Do two parallel tool calls."
)
print(str(response))
Added user message to memory: Tell me both the risk factors and tailwinds for Uber? Do two parallel tool calls. === Calling Function === Calling function: uber_10k with args: {"input": "What are the risk factors for Uber in 2021?"} === Calling Function === Calling function: uber_10k with args: {"input": "What are the tailwinds for Uber in 2021?"} assistant: Based on the information provided, here are the risk factors for Uber in 2021: 1. Failure to offer or develop autonomous vehicle technologies, which could result in inferior performance or safety concerns compared to competitors. 2. Dependence on high-quality personnel and the potential impact of attrition or unsuccessful succession planning on the business. 3. Security or data privacy breaches, unauthorized access, or destruction of proprietary, employee, or user data. 4. Cyberattacks, such as malware, ransomware, viruses, spamming, and phishing attacks, which could harm the company's reputation and operations. 5. Climate change risks, including physical and transitional risks, that may adversely impact the business if not managed effectively. 6. Reliance on third parties to maintain open marketplaces for distributing products and providing software, which could negatively affect the business if interfered with. 7. The need for additional capital to support business growth, which may not be available on reasonable terms or at all. 8. Difficulties in identifying, acquiring, and integrating suitable businesses, which could harm operating results and prospects. 9. Legal and regulatory risks, including extensive government regulation and oversight related to payment and financial services. 10. Intellectual property risks, such as the inability to protect intellectual property or claims of misappropriation by third parties. 11. Volatility in the market price of common stock, which could result in steep declines and loss of investment for shareholders. 12. Economic risks related to the COVID-19 pandemic, which has adversely impacted and could continue to adversely impact the business, financial condition, and results of operations. 13. The potential reclassification of Drivers as employees, workers, or quasi-employees, which could result in material costs associated with defending, settling, or resolving lawsuits and demands for arbitration. On the other hand, here are some tailwinds for Uber in 2021: 1. Launch of Uber One, a single cross-platform membership program in the United States, which offers discounts, special pricing, priority service, and exclusive perks across rides, delivery, and grocery offerings. 2. Introduction of a "Super App" view on iOS