← Back to Cookbook
Adaptive RAG
Details
File: third_party/LlamaIndex/Adaptive_RAG.ipynb
Type: Jupyter Notebook
Use Cases: Adaptive RAG, RAG
Integrations: Llamaindex
Content
Notebook content (JSON format):
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "8034fdda-462b-4751-aadb-ecac3b1d312c",
"metadata": {},
"source": [
"# Adaptive RAG with LlamaIndex\n",
"\n",
"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](https://arxiv.org/abs/2403.14403) proposes an approach to handle complex queries and simple queries.\n",
"\n",
"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.\n",
"\n",
"Our approach will involve using `RouterQueryEngine` and `FunctionCalling` capabilities of `MistralAI` to call different tools or indices based on the query's complexity.\n",
"\n",
"- **Complex Queries:** These will leverage multiple tools that require context from several documents.\n",
"- **Simple Queries:** These will utilize a single tool that requires context from a single document or directly use an LLM to provide an answer.\n",
"\n",
"Following are the steps we follow here:\n",
"\n",
"1. Download Data.\n",
"2. Load Data.\n",
"3. Create indices for 3 documents.\n",
"4. Create query engines with documents and LLM.\n",
"5. Initialize a `FunctionCallingAgentWorker` for complex queries.\n",
"6. Create tools.\n",
"7. Create `RouterQueryEngine` - To route queries based on its complexity.\n",
"8. Querying."
]
},
{
"cell_type": "markdown",
"id": "6e1b3314-e5af-4348-8a1b-661f0de3f48b",
"metadata": {},
"source": [
"### Installation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "310da4da-6487-446d-bd99-67e685bf6df5",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"!pip install llama-index\n",
"!pip install llama-index-llms-mistralai\n",
"!pip install llama-index-embeddings-mistralai"
]
},
{
"cell_type": "markdown",
"id": "39ae2bab-3116-4aab-a3c4-7670766f9625",
"metadata": {},
"source": [
"### Setup API Key"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "dbaf6798-95bb-47b0-9086-cdf0547ef00b",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRAL API KEY>'"
]
},
{
"cell_type": "markdown",
"id": "97b6c5be-6b72-4e4d-87ce-bc0a4e554c78",
"metadata": {},
"source": [
"### Setup LLM and Embedding Model"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b26e4ba8-e0ae-401f-9cc0-39544eb0d39f",
"metadata": {},
"outputs": [],
"source": [
"import nest_asyncio\n",
"\n",
"nest_asyncio.apply()"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "8a00879b-fa83-41b9-a00a-2ac02888fa11",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
"from llama_index.llms.mistralai import MistralAI\n",
"from llama_index.embeddings.mistralai import MistralAIEmbedding\n",
"from llama_index.core import Settings\n",
"\n",
"from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
"from llama_index.core.query_engine.router_query_engine import RouterQueryEngine\n",
"from llama_index.core.selectors.llm_selectors import LLMSingleSelector"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "c4d1dab9-bf0f-4e71-8763-e5dba29e4d1a",
"metadata": {},
"outputs": [],
"source": [
"# Note: Only `mistral-large-latest` supports function calling\n",
"llm = MistralAI(model='mistral-large-latest') \n",
"embed_model = MistralAIEmbedding()\n",
"\n",
"Settings.llm = llm\n",
"Settings.embed_model = embed_model"
]
},
{
"cell_type": "markdown",
"id": "0a6fc41a-7371-4c78-9c92-2bcab1273d98",
"metadata": {},
"source": [
"### Logging"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "ad61412c-bd14-4148-913a-0e7ae0da30c9",
"metadata": {},
"outputs": [],
"source": [
"# NOTE: This is ONLY necessary in jupyter notebook.\n",
"# Details: Jupyter runs an event-loop behind the scenes.\n",
"# This results in nested event-loops when we start an event-loop to make async queries.\n",
"# This is normally not allowed, we use nest_asyncio to allow it for convenience.\n",
"import nest_asyncio\n",
"\n",
"nest_asyncio.apply()\n",
"\n",
"import logging\n",
"import sys\n",
"\n",
"# Set up the root logger\n",
"logger = logging.getLogger()\n",
"logger.setLevel(logging.INFO) # Set logger level to INFO\n",
"\n",
"# Clear out any existing handlers\n",
"logger.handlers = []\n",
"\n",
"# Set up the StreamHandler to output to sys.stdout (Colab's output)\n",
"handler = logging.StreamHandler(sys.stdout)\n",
"handler.setLevel(logging.INFO) # Set handler level to INFO\n",
"\n",
"# Add the handler to the logger\n",
"logger.addHandler(handler)\n",
"\n",
"from IPython.display import display, HTML"
]
},
{
"cell_type": "markdown",
"id": "7d54a084-640d-45a2-babc-89bb12ee101b",
"metadata": {},
"source": [
"### Download Data\n",
"\n",
"We will download Lyft's 10k SEC filings for the years 2020, 2021, and 2022."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "d69a892d-7caa-44a6-a053-f363d058755d",
"metadata": {},
"outputs": [],
"source": [
"!wget \"https://www.dropbox.com/scl/fi/ywc29qvt66s8i97h1taci/lyft-10k-2020.pdf?rlkey=d7bru2jno7398imeirn09fey5&dl=0\" -q -O ./lyft_10k_2020.pdf\n",
"!wget \"https://www.dropbox.com/scl/fi/lpmmki7a9a14s1l5ef7ep/lyft-10k-2021.pdf?rlkey=ud5cwlfotrii6r5jjag1o3hvm&dl=0\" -q -O ./lyft_10k_2021.pdf\n",
"!wget \"https://www.dropbox.com/scl/fi/iffbbnbw9h7shqnnot5es/lyft-10k-2022.pdf?rlkey=grkdgxcrib60oegtp4jn8hpl8&dl=0\" -q -O ./lyft_10k_2022.pdf"
]
},
{
"cell_type": "markdown",
"id": "3491f538-fd9c-4875-bff5-da7b34a275f5",
"metadata": {},
"source": [
"### Load Data"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "f4d4e4cc-388c-43cf-a3a2-9af16292f510",
"metadata": {},
"outputs": [],
"source": [
"# Lyft 2020 docs\n",
"lyft_2020_docs = SimpleDirectoryReader(input_files=[\"./lyft_10k_2020.pdf\"]).load_data()\n",
"\n",
"# Lyft 2021 docs\n",
"lyft_2021_docs = SimpleDirectoryReader(input_files=[\"./lyft_10k_2021.pdf\"]).load_data()\n",
"\n",
"# Lyft 2022 docs\n",
"lyft_2022_docs = SimpleDirectoryReader(input_files=[\"./lyft_10k_2022.pdf\"]).load_data()"
]
},
{
"cell_type": "markdown",
"id": "40b3afe2-1364-4829-b378-30f150e83cbc",
"metadata": {},
"source": [
"### Create Indicies"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "cd4d2310-19fa-4f97-8488-20328069183e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n"
]
}
],
"source": [
"# Index on Lyft 2020 Document\n",
"lyft_2020_index = VectorStoreIndex.from_documents(lyft_2020_docs)\n",
"\n",
"# Index on Lyft 2021 Document\n",
"lyft_2021_index = VectorStoreIndex.from_documents(lyft_2021_docs)\n",
"\n",
"# Index on Lyft 2022 Document\n",
"lyft_2022_index = VectorStoreIndex.from_documents(lyft_2022_docs)"
]
},
{
"cell_type": "markdown",
"id": "b23c5588-52e9-4de0-91d2-192fc25a5ae9",
"metadata": {},
"source": [
"### Create Query Engines"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "079239e7-1ccb-4c53-b045-75286bb3dd84",
"metadata": {},
"outputs": [],
"source": [
"# Query Engine on Lyft 2020 Docs Index\n",
"lyft_2020_query_engine = lyft_2020_index.as_query_engine(similarity_top_k=5)\n",
"\n",
"# Query Engine on Lyft 2021 Docs Index\n",
"lyft_2021_query_engine = lyft_2021_index.as_query_engine(similarity_top_k=5)\n",
"\n",
"# Query Engine on Lyft 2022 Docs Index\n",
"lyft_2022_query_engine = lyft_2022_index.as_query_engine(similarity_top_k=5)"
]
},
{
"cell_type": "markdown",
"id": "93dd606d-0eb5-473e-aef9-5a38f308bb41",
"metadata": {},
"source": [
"Query Engine for LLM. With this we will use LLM to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "fa7a5d6a-215a-4754-8673-b3f4b39a7875",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.query_engine import CustomQueryEngine\n",
"\n",
"class LLMQueryEngine(CustomQueryEngine):\n",
" \"\"\"RAG String Query Engine.\"\"\"\n",
"\n",
" llm: llm\n",
"\n",
" def custom_query(self, query_str: str):\n",
"\n",
" response = self.llm.complete(query_str)\n",
"\n",
" return str(response)\n",
"\n",
"llm_query_engine = LLMQueryEngine(llm=llm)"
]
},
{
"cell_type": "markdown",
"id": "409f5c5d-c846-49ab-92ab-d4c3217ce4f6",
"metadata": {},
"source": [
"### Initialize a `FunctionCallingAgentWorker`"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "aad45c44-8218-41f7-87f2-36aae23bf65e",
"metadata": {},
"outputs": [],
"source": [
"# These tools are used to answer complex queries involving multiple documents.\n",
"query_engine_tools = [\n",
" QueryEngineTool(\n",
" query_engine=lyft_2020_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2020_10k_form\",\n",
" description=\"Annual report of Lyft's financial activities in 2020\",\n",
" ),\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=lyft_2021_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2021_10k_form\",\n",
" description=\"Annual report of Lyft's financial activities in 2021\",\n",
" ),\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=lyft_2022_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2022_10k_form\",\n",
" description=\"Annual report of Lyft's financial activities in 2022\",\n",
" ),\n",
" )\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "18dd9733-ce58-4edc-9a6f-f8c38dd29bc0",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.agent import FunctionCallingAgentWorker\n",
"from llama_index.core.agent import AgentRunner\n",
"\n",
"agent_worker = FunctionCallingAgentWorker.from_tools(\n",
" query_engine_tools,\n",
" llm=llm,\n",
" verbose=True,\n",
" allow_parallel_tool_calls=True,\n",
")\n",
"agent = AgentRunner(agent_worker)"
]
},
{
"cell_type": "markdown",
"id": "504e13c0-487d-4713-8bb4-10cac88e8ff6",
"metadata": {},
"source": [
"### Create Tools\n",
"\n",
"We will create tools using the `QueryEngines`, and `FunctionCallingAgentWorker` created earlier."
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "7b3b5737-5d64-4f6a-833a-b756651285b3",
"metadata": {},
"outputs": [],
"source": [
"query_engine_tools = [\n",
" QueryEngineTool(\n",
" query_engine=lyft_2020_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2020_10k_form\",\n",
" description=\"Queries related to only 2020 Lyft's financial activities.\",\n",
" ),\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=lyft_2021_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2021_10k_form\",\n",
" description=\"Queries related to only 2021 Lyft's financial activities.\",\n",
" ),\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=lyft_2022_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2022_10k_form\",\n",
" description=\"Queries related to only 2022 Lyft's financial activities.\",\n",
" ),\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=agent,\n",
" metadata=ToolMetadata(\n",
" name=\"lyft_2020_2021_2022_10k_form\",\n",
" description=(\n",
" \"Useful for queries that span multiple years from 2020 to 2022 for Lyft's financial activities.\"\n",
"\n",
" )\n",
" )\n",
" ),\n",
" QueryEngineTool(\n",
" query_engine=llm_query_engine,\n",
" metadata=ToolMetadata(\n",
" name=\"general_queries\",\n",
" description=(\n",
" \"Provides information about general queries other than lyft.\"\n",
" )\n",
" )\n",
" )\n",
"]"
]
},
{
"cell_type": "markdown",
"id": "1362a167-a45d-4171-a770-c52bba98da9a",
"metadata": {},
"source": [
"### Create RouterQueryEngine\n",
"\n",
"`RouterQueryEngine` will route user queries to select one of the tools based on the complexity of the query."
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "454b01d7-b70b-4f10-8b7d-5c533d504a68",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.query_engine import RouterQueryEngine\n",
"from llama_index.core.selectors import LLMSingleSelector\n",
"\n",
"query_engine = RouterQueryEngine(\n",
" selector=LLMSingleSelector.from_defaults(),\n",
" query_engine_tools=query_engine_tools,\n",
" verbose = True\n",
")"
]
},
{
"cell_type": "markdown",
"id": "e1debe82-1016-4ade-ba3f-284f0c872b79",
"metadata": {},
"source": [
"### Querying"
]
},
{
"cell_type": "markdown",
"id": "33f0ee2c-2e7b-4301-bb0d-0fed29b75e88",
"metadata": {},
"source": [
"#### Simple Queries:"
]
},
{
"cell_type": "markdown",
"id": "e26f0f65-129f-4774-b881-68f370fa9771",
"metadata": {},
"source": [
"##### Query: What is the capital of France?\n",
"\n",
"You can see that it used LLM tool since it is a general query."
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "f1f22b16-2067-4e68-8672-f217e85e3873",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mHTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">The capital of France is Paris. Known as the \"City of Light,\" Paris is famous for its iconic landmarks like the Eiffel Tower, Louvre Museum, Notre-Dame Cathedral, and more. It is a global center for art, fashion, gastronomy, and culture.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What is the capital of France?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"cell_type": "markdown",
"id": "21dfee51-b28c-40b4-91fd-f84137f8b760",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2022?\n",
"\n",
"You can see that it used lyft_2022 tool to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "aeccdd6c-9bf9-44a1-b788-8f3a44f7da57",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mHTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">In 2022, Lyft made focused investments in research and development, though the specific details of these investments are not provided in the context. These investments are part of their ongoing commitment to launch new innovations on their platform. Additionally, Lyft completed multiple strategic acquisitions, one of which is PBSC Urban Solutions. This acquisition added over 100,000 bikes to their bikeshare systems across 46 markets in 15 countries.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2022?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"cell_type": "markdown",
"id": "bd3c8ea7-8db3-40f3-b8a6-c548037edefe",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2021?\n",
"\n",
"You can see that it used lyft_2021 tool to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "9dbc0894-ecfb-43cd-86f4-566d2e21d602",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mHTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">In 2021, Lyft's research and development expenses increased slightly by $2.8 million. This increase was primarily due to a $51.6 million rise in stock-based compensation. However, there was also a $25.4 million benefit from a restructuring event in the second quarter of 2020, which included a stock-based compensation benefit and severance and benefits costs that did not recur in 2021. These increases were offset by a $37.5 million decrease in personnel-related costs and a $4.6 million decrease in autonomous vehicle research costs.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2021?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"cell_type": "markdown",
"id": "6a2f002d-7a30-4634-96b6-da32cc4b02e1",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2020?\n",
"\n",
"You can see that it used lyft_2020 tool to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "c2189dfe-fe62-4b93-bdf6-2f8e422e6672",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mHTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">In 2020, Lyft continued to invest heavily in research and development as part of its mission to revolutionize transportation. The company made focused investments to launch new innovations on its platform. One notable example is the acquisition of Flexdrive, LLC, a longstanding partner in Lyft's Express Drive program. This acquisition is expected to contribute to the growth of Lyft's business and help expand the range of its use cases. Additionally, Lyft invested in the expansion of its network of Light Vehicles and autonomous open platform technology to remain at the forefront of transportation innovation. Despite the impact of COVID-19, Lyft plans to continue investing in the future, both organically and through acquisitions of complementary businesses.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2020?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"cell_type": "markdown",
"id": "077ebe1a-bbc2-40ff-a51e-ae57e48d432a",
"metadata": {},
"source": [
"#### Complex Queries\n",
"\n",
"Let's test queries that requires multiple tools."
]
},
{
"cell_type": "markdown",
"id": "3f1a5602-aba6-4171-b90f-9afce7d4911c",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2022 vs 2020?\n",
"\n",
"You can see that it used lyft_2020 and lyft_2022 tools with `FunctionCallingAgent` to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "9a966798-4cbe-4ce2-bbac-7af8632cacb2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mAdded user message to memory: What did Lyft do in R&D in 2022 vs 2020?\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2020_10k_form with args: {\"input\": \"R&D expenses\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2022_10k_form with args: {\"input\": \"R&D expenses\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">assistant: In 2020, Lyft's Research and Development expenses were $25,376 million, after accounting for a stock-based compensation benefit of $37,082 million and severance and other employee costs of $11,706 million. In 2022, these expenses decreased to $856,777 thousand, a decrease of $55.2 million, or 6%. This decrease was primarily due to a reduction in personnel-related costs and stock-based compensation, driven by reduced headcount following a transaction with Woven Planet in Q3 2021. There were also decreases in Level 5 costs, web hosting fees, autonomous vehicle research costs, and consulting and advisory costs. However, these decreases were partially offset by restructuring costs related to an event in Q4 2022.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2022 vs 2020?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"cell_type": "markdown",
"id": "689d367c-2a35-462b-9b61-7fbf1e93fcda",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2021 vs 2020?\n",
"\n",
"You can see that it used lyft_2020 and lyft_2021 tools with `FunctionCallingAgent` to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "3235100b-d552-4a8d-9203-19889ce28e4c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mAdded user message to memory: What did Lyft do in R&D in 2020 vs 2021?\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2020_10k_form with args: {\"input\": \"R&D activities\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2021_10k_form with args: {\"input\": \"R&D activities\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">assistant: In 2020, Lyft made significant investments in research and development to launch new innovations on their platform. They acquired Flexdrive, LLC, one of their longstanding partners in the Express Drive program, which is expected to contribute to the growth of their business and help expand the range of their use cases. They also invested in the expansion of their network of Light Vehicles and autonomous open platform technology.\n",
"\n",
"In 2021, Lyft's R&D activities primarily consisted of personnel-related compensation costs and facilities costs, including costs related to autonomous vehicle technology initiatives. However, following a transaction with Woven Planet, a subsidiary of Toyota Motor Corporation, certain costs related to Lyft's prior initiative to develop self-driving systems were eliminated beginning in the third quarter of 2021. This transaction involved the divestiture of certain assets related to Lyft's self-driving vehicle division, Level 5.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2020 vs 2021?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "e1189418-c0c7-474f-981d-59160b999583",
"metadata": {},
"source": [
"##### Query: What did Lyft do in R&D in 2022 vs 2021 vs 2020?\n",
"\n",
"You can see that it used lyft_2020, lyft_2021 and lyft_2022 tools with `FunctionCallingAgent` to answer the query."
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "e14adac5-a2e8-4224-beea-ba8ae7b3037d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"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..\n",
"\u001b[1;3;38;5;200mSelecting 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..\n",
"\u001b[0mAdded user message to memory: What did Lyft do in R&D in 2022 vs 2021 vs 2020?\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2020_10k_form with args: {\"input\": \"R&D\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2021_10k_form with args: {\"input\": \"R&D\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"=== Calling Function ===\n",
"Calling function: lyft_2022_10k_form with args: {\"input\": \"R&D\"}\n",
"HTTP Request: POST https://api.mistral.ai/v1/embeddings \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
"HTTP Request: POST https://api.mistral.ai/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
]
},
{
"data": {
"text/html": [
"<p style=\"font-size:20px\">assistant: In 2020, Lyft made substantial investments in research and development, with a focus on launching new innovations on their platform. They held 342 issued U.S. patents and had 446 U.S. patent applications pending, along with 59 issued patents in foreign jurisdictions and 223 applications pending in foreign jurisdictions.\n",
"\n",
"In 2021, Lyft's R&D expenses primarily consisted of personnel-related compensation costs and facilities costs. They also had costs related to autonomous vehicle technology initiatives. However, on July 13, 2021, Lyft completed a transaction with Woven Planet, a subsidiary of Toyota Motor Corporation, for the divestiture of certain assets related to their self-driving vehicle division, Level 5. As a result, certain costs related to their prior initiative to develop self-driving systems were eliminated beginning in the third quarter of 2021.\n",
"\n",
"In 2022, Lyft's R&D expenses decreased by $55.2 million, or 6%, compared to the previous year. This decrease was primarily due to a reduction in personnel-related costs and stock-based compensation, driven by reduced headcount following the transaction with Woven Planet in the third quarter of 2021. There were also decreases in Level 5 costs, web hosting fees, autonomous vehicle research costs, and consulting and advisory costs. However, these decreases were partially offset by restructuring costs in the fourth quarter of 2022, which included impairment costs of operating lease right-of-use assets, severance and benefits costs, and stock-based compensation.</p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"response = query_engine.query(\"What did Lyft do in R&D in 2022 vs 2021 vs 2020?\")\n",
"display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}