Mistral AI Adaptive RAG

Adaptive RAG

mistral-cookbookLlamaIndexthird_party

alph-notebooks/mistral-cookbook / Adaptive_RAG.ipynb

Export

Run Notebooks

Contents

No cells yet

Add cells to see them here

Adaptive RAG with LlamaIndex

User queries in general can be complex queries, simple queries. One don't always need complex RAG system even to handle simple queries. Adaptive RAG proposes an approach to handle complex queries and simple queries.

In this notebook, we will implement an approach similar to Adaptive RAG, which differentiates between handling complex and simple queries. We'll focus on Lyft's 10k SEC filings for the years 2020, 2021, and 2022.

Our approach will involve using RouterQueryEngine and FunctionCalling capabilities of MistralAI to call different tools or indices based on the query's complexity.

Complex Queries: These will leverage multiple tools that require context from several documents.
Simple Queries: These will utilize a single tool that requires context from a single document or directly use an LLM to provide an answer.

Following are the steps we follow here:

Download Data.
Load Data.
Create indices for 3 documents.
Create query engines with documents and LLM.
Initialize a FunctionCallingAgentWorker for complex queries.
Create tools.
Create RouterQueryEngine - To route queries based on its complexity.
Querying.

Installation

[ ]

Setup API Key

[1]

Setup LLM and Embedding Model

[2]

[3]

[4]

Logging

[5]

Download Data

We will download Lyft's 10k SEC filings for the years 2020, 2021, and 2022.

[6]

Load Data

[7]

Create Indicies

[8]

HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"

Create Query Engines

[9]

Query Engine for LLM. With this we will use LLM to answer the query.

[10]

Initialize a `FunctionCallingAgentWorker`

[11]

[12]

Create Tools

We will create tools using the QueryEngines, and FunctionCallingAgentWorker created earlier.

[13]

Create RouterQueryEngine

RouterQueryEngine will route user queries to select one of the tools based on the complexity of the query.

[14]

Querying

Simple Queries:

Query: What is the capital of France?

You can see that it used LLM tool since it is a general query.

[15]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 4: This option is the most relevant as it mentions 'general queries other than Lyft'. The question about the capital of France is not related to Lyft's financial activities..
Selecting query engine 4: This option is the most relevant as it mentions 'general queries other than Lyft'. The question about the capital of France is not related to Lyft's financial activities..
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Query: What did Lyft do in R&D in 2022?

You can see that it used lyft_2022 tool to answer the query.

[16]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 2: This choice is most relevant to the question as it pertains to Lyft's financial activities in 2022, which could include information about their Research and Development (R&D) expenditures for that year..
Selecting query engine 2: This choice is most relevant to the question as it pertains to Lyft's financial activities in 2022, which could include information about their Research and Development (R&D) expenditures for that year..
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Query: What did Lyft do in R&D in 2021?

You can see that it used lyft_2021 tool to answer the query.

[17]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 1: This choice is most relevant to the question as it pertains to queries related to only 2021 Lyft's financial activities, which could potentially include information about Lyft's R&D activities during that year..
Selecting query engine 1: This choice is most relevant to the question as it pertains to queries related to only 2021 Lyft's financial activities, which could potentially include information about Lyft's R&D activities during that year..
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Query: What did Lyft do in R&D in 2020?

You can see that it used lyft_2020 tool to answer the query.

[18]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 0: This choice is most relevant to the question as it pertains to queries related to only 2020 Lyft's financial activities, which could potentially include information about Lyft's R&D activities during that year..
Selecting query engine 0: This choice is most relevant to the question as it pertains to queries related to only 2020 Lyft's financial activities, which could potentially include information about Lyft's R&D activities during that year..
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Complex Queries

Let's test queries that requires multiple tools.

Query: What did Lyft do in R&D in 2022 vs 2020?

You can see that it used lyft_2020 and lyft_2022 tools with FunctionCallingAgent to answer the query.

[19]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would allow for a comparison of Lyft's R&D activities in 2022 versus 2020..
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would allow for a comparison of Lyft's R&D activities in 2022 versus 2020..
Added user message to memory: What did Lyft do in R&D in 2022 vs 2020?
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2020_10k_form with args: {"input": "R&D expenses"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2022_10k_form with args: {"input": "R&D expenses"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Query: What did Lyft do in R&D in 2021 vs 2020?

You can see that it used lyft_2020 and lyft_2021 tools with FunctionCallingAgent to answer the query.

[20]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would cover the comparison of Lyft's R&D activities in 2020 vs 2021..
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would cover the comparison of Lyft's R&D activities in 2020 vs 2021..
Added user message to memory: What did Lyft do in R&D in 2020 vs 2021?
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2020_10k_form with args: {"input": "R&D activities"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2021_10k_form with args: {"input": "R&D activities"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"

Query: What did Lyft do in R&D in 2022 vs 2021 vs 2020?

You can see that it used lyft_2020, lyft_2021 and lyft_2022 tools with FunctionCallingAgent to answer the query.

[21]

HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would allow for a comparison of Lyft's R&D activities across these years..
Selecting query engine 3: This choice is most relevant as it spans multiple years from 2020 to 2022, which would allow for a comparison of Lyft's R&D activities across these years..
Added user message to memory: What did Lyft do in R&D in 2022 vs 2021 vs 2020?
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2020_10k_form with args: {"input": "R&D"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2021_10k_form with args: {"input": "R&D"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
=== Calling Function ===
Calling function: lyft_2022_10k_form with args: {"input": "R&D"}
HTTP Request: POST https://api.mistral.ai/v1/embeddings "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"
HTTP Request: POST https://api.mistral.ai/v1/chat/completions "HTTP/1.1 200 OK"