Build a ReAct Agents with Mistral AI and LlamaIndex

RAGLlamaindex+1

This notebook shows you how to use ReAct Agent and FunctionCalling Agent over defined tools and RAG pipeline with MistralAI LLM.

Installation#

If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

!pip install llama-index
!pip install llama-index-llms-mistralai
!pip install llama-index-embeddings-mistralai

Setup API Key#

import os
os.environ['MISTRAL_API_KEY'] = 'YOUR MISTRAL API KEY'
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")

With FunctionCalling Agent#

Here we initialize a simple FunctionCalling agent with calculator functions.

from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=False,
)
agent = AgentRunner(agent_worker)

Chat#

response = agent.chat("What is (121 + 2) * 5?")
print(str(response))
# inspect sources
print(response.sources)

Async Chat#

Also let's re-enable parallel function calling so that we can call two multiply operations simultaneously.

# enable parallel function calling
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=True,
)
agent = AgentRunner(agent_worker)
response = await agent.achat("What is (121 * 3) + (5 * 8)?")
print(str(response))

With ReAct Agent#

from llama_index.core.agent import ReActAgent

agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)

response = agent.chat("What is (121 * 3) + (5 * 8)?")
print(str(response))

Agent over RAG Pipeline#

Build a Mistral FunctionCalling 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'
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."
        ),
    ),
)

With FunctionCalling Agent#

from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

agent_worker = FunctionCallingAgentWorker.from_tools(
    [query_engine_tool], llm=llm, verbose=True
)
agent = AgentRunner(agent_worker)
response = agent.chat(
    "What are the risk factors for Uber in 2021?"
)
print(str(response))

With ReAct Agent#

from llama_index.core.agent import ReActAgent

agent = ReActAgent.from_tools([query_engine_tool], llm=llm, verbose=True)

response = agent.chat("What are the risk factors for Uber in 2021?")
print(str(response))