recipes/llama_api_providers/llama3_cookbook_groq.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "09211e76-286f-4b12-acd7-cfb082dc2d66",
   "metadata": {},
   "source": [
    "# Llama 3 Cookbook with LlamaIndex and Groq\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-recipes/blob/main/recipes/llama_api_providers/llama3_cookbook_groq.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "Meta developed and released the Meta [Llama 3](https://ai.meta.com/blog/meta-llama-3/) family of large language models (LLMs), a collection of pretrained and instruction tuned generative text models in 8 and 70B sizes. The Llama 3 instruction tuned models are optimized for dialogue use cases and outperform many of the available open source chat models on common industry benchmarks.\n",
    "\n",
    "In this notebook, we demonstrate how to use Llama 3 with LlamaIndex for a comprehensive set of use cases. \n",
    "1. Basic completion / chat \n",
    "2. Basic RAG (Vector Search, Summarization)\n",
    "3. Advanced RAG (Routing)\n",
    "4. Text-to-SQL \n",
    "5. Structured Data Extraction\n",
    "6. Chat Engine + Memory\n",
    "7. Agents\n",
    "\n",
    "\n",
    "We use Llama3-8B and Llama3-70B through [Groq](https://groq.com) - you can sign up there to get a free trial API key."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de2901c0-e20d-48e5-9385-dbca2258c564",
   "metadata": {},
   "source": [
    "## Installation and Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf643ac-b025-4812-aaed-f8f85d1ba505",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index\n",
    "!pip install llama-index-llms-groq\n",
    "!pip install llama-index-embeddings-huggingface\n",
    "!pip install llama-parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641fa5c8-d63e-47f8-b5bc-ebf994f6e314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1714ea83-6cd4-44bb-b53f-4499126c3809",
   "metadata": {},
   "source": [
    "### Setup LLM using Groq\n",
    "\n",
    "To use [Groq](https://groq.com), you need to make sure that `GROQ_API_KEY` is specified as an environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d46440c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"GROQ_API_KEY\"] = \"YOUR_GROQ_API_KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5256970-eba4-499a-b438-8766a290a61a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.groq import Groq\n",
    "\n",
    "llm = Groq(model=\"llama3-8b-8192\")\n",
    "llm_70b = Groq(model=\"llama3-70b-8192\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c3f154-d345-465d-8eed-63b99adbd3ca",
   "metadata": {},
   "source": [
    "### Setup Embedding Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cda736d-e414-44e3-8c15-6be49f5f0282",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3625cf29-7c56-475a-8efd-fbe8ffce194d",
   "metadata": {},
   "source": [
    "### Define Global Settings Configuration\n",
    "\n",
    "In LlamaIndex, you can define global settings so you don't have to pass the LLM / embedding model objects everywhere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be3565d1-cc5b-4149-ad5a-7be8f7818e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42449b68-47f5-40cf-9207-191307b25e8e",
   "metadata": {},
   "source": [
    "### Download Data\n",
    "\n",
    "Here you'll download data that's used in section 2 and onwards.\n",
    "\n",
    "We'll download some articles on Kendrick, Drake, and their beef (as of May 2024)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b18640-cdfa-42c1-ab53-115983c1fdc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "!wget \"https://www.dropbox.com/scl/fi/t1soxfjdp0v44an6sdymd/drake_kendrick_beef.pdf?rlkey=u9546ymb7fj8lk2v64r6p5r5k&st=wjzzrgil&dl=1\" -O data/drake_kendrick_beef.pdf\n",
    "!wget \"https://www.dropbox.com/scl/fi/nts3n64s6kymner2jppd6/drake.pdf?rlkey=hksirpqwzlzqoejn55zemk6ld&st=mohyfyh4&dl=1\" -O data/drake.pdf\n",
    "!wget \"https://www.dropbox.com/scl/fi/8ax2vnoebhmy44bes2n1d/kendrick.pdf?rlkey=fhxvn94t5amdqcv9vshifd3hj&st=dxdtytn6&dl=1\" -O data/kendrick.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9edee491-05f8-4fbb-9394-baa82f1e5087",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "We load data using LlamaParse by default, but you can also choose to opt for our free pypdf reader (in SimpleDirectoryReader by default) if you don't have an account! \n",
    "\n",
    "1. LlamaParse: Signup for an account here: cloud.llamaindex.ai. You get 1k free pages a day, and paid plan is 7k free pages + 0.3c per additional page. LlamaParse is a good option if you want to parse complex documents, like PDFs with charts, tables, and more. \n",
    "\n",
    "2. Default PDF Parser (In `SimpleDirectoryReader`). If you don't want to signup for an account / use a PDF service, just use the default PyPDF reader bundled in our file loader. It's a good choice for getting started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b648635a-2672-407f-bae6-01660e5426d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment this code if you want to use LlamaParse\n",
    "# from llama_parse import LlamaParse\n",
    "\n",
    "# docs_kendrick = LlamaParse(result_type=\"text\").load_data(\"./data/kendrick.pdf\")\n",
    "# docs_drake = LlamaParse(result_type=\"text\").load_data(\"./data/drake.pdf\")\n",
    "# docs_both = LlamaParse(result_type=\"text\").load_data(\n",
    "#     \"./data/drake_kendrick_beef.pdf\"\n",
    "# )\n",
    "\n",
    "# Uncomment this code if you want to use SimpleDirectoryReader / default PDF Parser\n",
    "# from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# docs_kendrick = SimpleDirectoryReader(input_files=[\"data/kendrick.pdf\"]).load_data()\n",
    "# docs_drake = SimpleDirectoryReader(input_files=[\"data/drake.pdf\"]).load_data()\n",
    "# docs_both = SimpleDirectoryReader(input_files=[\"data/drake_kendrick_beef.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "071a8f44-2765-4d57-b8da-15d3c718874d",
   "metadata": {},
   "source": [
    "## 1. Basic Completion and Chat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0b1ace8-32fb-46b2-a065-8817ddc0310b",
   "metadata": {},
   "source": [
    "### Call complete with a prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2db43f9-74af-453c-9f83-8db0379c3302",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm.complete(\"do you like drake or kendrick better?\")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89326153-e2d2-4136-8193-fb27d20670c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_response = llm.stream_complete(\n",
    "    \"you're a drake fan. tell me why you like drake more than kendrick\"\n",
    ")\n",
    "\n",
    "for t in stream_response:\n",
    "    print(t.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4558339-c8a1-4d26-a430-eb71768b5351",
   "metadata": {},
   "source": [
    "### Call chat with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f393031-f743-4a28-a122-71817e3fbd1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are Kendrick.\"),\n",
    "    ChatMessage(role=\"user\", content=\"Write a verse.\"),\n",
    "]\n",
    "response = llm.chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9551fc-0efc-4671-bc57-339121004c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a67a33d-fe7d-4381-983f-ca3a6945995d",
   "metadata": {},
   "source": [
    "## 2. Basic RAG (Vector Search, Summarization)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c104a0c5-e43b-475b-9fa6-186906c1f327",
   "metadata": {},
   "source": [
    "### Basic RAG (Vector Search)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216787b7-e40a-43fc-a4ca-c43cb798ce9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(docs_both)\n",
    "query_engine = index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a854e9d3-70f1-4927-a2f6-59e90c31f2f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"Tell me about family matters\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da796970-bc38-4cb4-9d32-ebd1b71d4bdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eff935b7-4f37-4758-8997-82fb0852e732",
   "metadata": {},
   "source": [
    "### Basic RAG (Summarization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfe72300-7a38-453e-b1f2-bc1c00a01ff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SummaryIndex\n",
    "\n",
    "summary_index = SummaryIndex.from_documents(docs_both)\n",
    "summary_engine = summary_index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178f1f12-51f7-4b45-9346-c16ed12b3b8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = summary_engine.query(\n",
    "    \"Given your assessment of this article, who won the beef?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8125382-d576-4b99-a0da-2fbb71a5b19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68918eb6-f1e6-460c-b1d5-fb49c3fed4b8",
   "metadata": {},
   "source": [
    "## 3. Advanced RAG (Routing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94fd7097-0287-4522-8e43-3e088291fa8a",
   "metadata": {},
   "source": [
    "### Build a Router that can choose whether to do vector search or summarization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3949dd41-e9a1-47f6-900f-4f987cad3f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
    "\n",
    "vector_tool = QueryEngineTool(\n",
    "    index.as_query_engine(),\n",
    "    metadata=ToolMetadata(\n",
    "        name=\"vector_search\",\n",
    "        description=\"Useful for searching for specific facts.\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "summary_tool = QueryEngineTool(\n",
    "    index.as_query_engine(response_mode=\"tree_summarize\"),\n",
    "    metadata=ToolMetadata(\n",
    "        name=\"summary\",\n",
    "        description=\"Useful for summarizing an entire document.\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d063d07b-c03e-4b26-8556-e3c058d2fd52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import RouterQueryEngine\n",
    "\n",
    "query_engine = RouterQueryEngine.from_defaults(\n",
    "    [vector_tool, summary_tool], select_multi=False, verbose=True, llm=llm_70b\n",
    ")\n",
    "\n",
    "response = query_engine.query(\n",
    "    \"Tell me about the song meet the grahams - why is it significant\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396aad75-5a71-4bd9-a760-7f13fe223079",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a795f0bc-e871-4580-8983-6fb27d421fc5",
   "metadata": {},
   "source": [
    "## 4. Text-to-SQL \n",
    "\n",
    "Here, we download and use a sample SQLite database with 11 tables, with various info about music, playlists, and customers. We will limit to a select few tables for this test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5096501-92c3-41af-a871-ade869d710fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip\" -O \"./data/chinook.zip\"\n",
    "!unzip \"./data/chinook.zip\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4db989e-c18d-4416-928e-7be4ead4d869",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    "    select,\n",
    "    column,\n",
    ")\n",
    "\n",
    "engine = create_engine(\"sqlite:///chinook.db\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf6ed233-0ea3-4d4f-8c33-5b6d558b89b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SQLDatabase\n",
    "\n",
    "sql_database = SQLDatabase(engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "debae423-1004-40f6-9356-e1c3add4d965",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.indices.struct_store import NLSQLTableQueryEngine\n",
    "\n",
    "query_engine = NLSQLTableQueryEngine(\n",
    "    sql_database=sql_database,\n",
    "    tables=[\"albums\", \"tracks\", \"artists\"],\n",
    "    llm=llm_70b,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a65ecd70-09c4-4872-b712-3a8235d03db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What are some albums?\")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c12b93ef-d6d1-4d15-9cb2-343070f72851",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What are some artists? Limit it to 5.\")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c243d38-c6ac-445c-b9d4-53a9ae013b7b",
   "metadata": {},
   "source": [
    "This last query should be a more complex join"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "553741c2-1050-445d-979a-ae2150ee3248",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\n",
    "    \"What are some tracks from the artist AC/DC? Limit it to 3\"\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "300689d7-9e67-4404-9898-27404ee6d4b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response.metadata[\"sql_query\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1419fe67-aa6a-47db-88cd-9bb251c15615",
   "metadata": {},
   "source": [
    "## 5. Structured Data Extraction\n",
    "\n",
    "An important use case for function calling is extracting structured objects. LlamaIndex provides an intuitive interface for this through `structured_predict` - simply define the target Pydantic class (can be nested), and given a prompt, we extract out the desired object.\n",
    "\n",
    "**NOTE**: Since there's no native function calling support with Llama3, the structured extraction is performed by prompting the LLM + output parsing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4432f35a-5f29-45e9-a928-32e6d77b158e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.groq import Groq\n",
    "from llama_index.core.prompts import PromptTemplate\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class Restaurant(BaseModel):\n",
    "    \"\"\"A restaurant with name, city, and cuisine.\"\"\"\n",
    "\n",
    "    name: str\n",
    "    city: str\n",
    "    cuisine: str\n",
    "\n",
    "\n",
    "llm = Groq(model=\"llama3-8b-8192\", pydantic_program_mode=\"llm\")\n",
    "prompt_tmpl = PromptTemplate(\n",
    "    \"Generate a restaurant in a given city {city_name}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c451f52-a051-4ba2-a683-0c1fd258d986",
   "metadata": {},
   "outputs": [],
   "source": [
    "restaurant_obj = llm.structured_predict(\n",
    "    Restaurant, prompt_tmpl, city_name=\"Miami\"\n",
    ")\n",
    "print(restaurant_obj)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "839018a9-b65f-4824-83f7-2e4e52b55c5d",
   "metadata": {},
   "source": [
    "## 6. Adding Chat History to RAG (Chat Engine)\n",
    "\n",
    "In this section we create a stateful chatbot from a RAG pipeline, with our chat engine abstraction.\n",
    "\n",
    "Unlike a stateless query engine, the chat engine maintains conversation history (through a memory module like buffer memory). It performs retrieval given a condensed question, and feeds the condensed question + context + chat history into the final LLM prompt.\n",
    "\n",
    "Related resource: https://docs.llamaindex.ai/en/stable/examples/chat_engine/chat_engine_condense_plus_context/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e56315-9513-4b32-bf9a-ce97c3ab52df",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.memory import ChatMemoryBuffer\n",
    "from llama_index.core.chat_engine import CondensePlusContextChatEngine\n",
    "\n",
    "memory = ChatMemoryBuffer.from_defaults(token_limit=3900)\n",
    "\n",
    "chat_engine = CondensePlusContextChatEngine.from_defaults(\n",
    "    index.as_retriever(),\n",
    "    memory=memory,\n",
    "    llm=llm,\n",
    "    context_prompt=(\n",
    "        \"You are a chatbot, able to have normal interactions, as well as talk\"\n",
    "        \" about the Kendrick and Drake beef.\"\n",
    "        \"Here are the relevant documents for the context:\\n\"\n",
    "        \"{context_str}\"\n",
    "        \"\\nInstruction: Use the previous chat history, or the context above, to interact and help the user.\"\n",
    "    ),\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24524d2-fdce-4237-8ecc-67f139302303",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = chat_engine.chat(\n",
    "    \"Tell me about the songs Drake released in the beef.\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a87a16-2864-4c48-95e7-a2103e119242",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = chat_engine.chat(\"What about Kendrick?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7fa07ed-58f0-445e-bbd3-4ad8bac6598e",
   "metadata": {},
   "source": [
    "## 7. Agents\n",
    "\n",
    "Here we build agents with Llama 3. We perform RAG over simple functions as well as the documents above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa98d735-5d43-413f-aab3-fc3adeed81b1",
   "metadata": {},
   "source": [
    "### Agents And Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb73a01f-8a2e-4dd6-91f8-710c92b81c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Sequence, List\n",
    "\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core.tools import BaseTool, FunctionTool\n",
    "from llama_index.core.agent import ReActAgent\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efbee832-9786-4551-93f2-01ee90fa0f4d",
   "metadata": {},
   "source": [
    "### Define Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2058b36-8053-4dc8-9218-c286702ecf66",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two integers and returns the result integer\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "def divide(a: int, b: int) -> int:\n",
    "    \"\"\"Divides two integers and returns the result integer\"\"\"\n",
    "    return a / b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "subtract_tool = FunctionTool.from_defaults(fn=subtract)\n",
    "divide_tool = FunctionTool.from_defaults(fn=divide)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22d7d4dc-e2ce-402c-9350-0e7010d0080c",
   "metadata": {},
   "source": [
    "### ReAct Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72a48053-e30d-4884-bcac-80752047d940",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReActAgent.from_tools(\n",
    "    [multiply_tool, add_tool, subtract_tool, divide_tool],\n",
    "    llm=llm_70b,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ada828a-3b05-4fc1-90e8-986c5607ae61",
   "metadata": {},
   "source": [
    "### Querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c0b1e56-d9f7-4615-a15a-c91fea1adb00",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.chat(\"What is (121 + 2) * 5?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ce45f6-bdd4-42aa-8f74-43a50f14094e",
   "metadata": {},
   "source": [
    "### ReAct Agent With RAG QueryEngine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97fce5f1-eacf-4ecc-9e83-072e74d3a2a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "\n",
    "from llama_index.core.tools import QueryEngineTool, ToolMetadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23963d00-e3d2-4ce1-9ac3-aa486bf4b1a5",
   "metadata": {},
   "source": [
    "### Create ReAct Agent using RAG QueryEngine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1844dbbd-477c-4c4d-bb18-2c2e16a75a50",
   "metadata": {},
   "source": [
    "This may take 4 minutes to run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66ab1e60-3374-4eb9-b7dc-c28db3b47c51",
   "metadata": {},
   "outputs": [],
   "source": [
    "drake_index = VectorStoreIndex.from_documents(docs_drake)\n",
    "drake_query_engine = drake_index.as_query_engine(similarity_top_k=3)\n",
    "\n",
    "kendrick_index = VectorStoreIndex.from_documents(docs_kendrick)\n",
    "kendrick_query_engine = kendrick_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e241fe9-f390-4be5-b3c4-da4f56db01ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "drake_tool = QueryEngineTool(\n",
    "    drake_index.as_query_engine(),\n",
    "    metadata=ToolMetadata(\n",
    "        name=\"drake_search\",\n",
    "        description=\"Useful for searching over Drake's life.\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "kendrick_tool = QueryEngineTool(\n",
    "    kendrick_index.as_query_engine(),\n",
    "    metadata=ToolMetadata(\n",
    "        name=\"kendrick_search\",\n",
    "        description=\"Useful for searching over Kendrick's life.\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "query_engine_tools = [drake_tool, kendrick_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b922feac-b221-4737-92c6-e63eeab4eab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReActAgent.from_tools(\n",
    "    query_engine_tools,\n",
    "    llm=llm_70b,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e38edc8-47f8-4f1a-ad87-bc3a9e31a65e",
   "metadata": {},
   "source": [
    "### Querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "035c2c8b-5a5e-4df0-a423-4c2d6054f457",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.chat(\"Tell me about how Kendrick and Drake grew up\")\n",
    "print(str(response))"
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}