[2' read]

View on GitHub ↗Open in Colab ↗

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))