llamaindex_cookbook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "09211e76-286f-4b12-acd7-cfb082dc2d66",
   "metadata": {},
   "source": [
    "# Llama3 Cookbook with Ollama and Replicate\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/cookbooks/llama3_cookbook_ollama_replicate.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 Llama3 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, Sub-Questions)\n",
    "4. Text-to-SQL \n",
    "5. Structured Data Extraction\n",
    "6. Agents\n",
    "\n",
    "\n",
    "We use Llama3-8B through Ollama, and Llama3-70B through Replicate. "
   ]
  },
  {
   "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-ollama\n",
    "!pip install llama-index-llms-replicate\n",
    "!pip install llama-index-embeddings-huggingface\n",
    "!pip install llama-parse\n",
    "!pip install replicate"
   ]
  },
  {
   "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 Ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5256970-eba4-499a-b438-8766a290a61a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ollama import Ollama\n",
    "\n",
    "llm = Ollama(model=\"llama3\", request_timeout=120.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "729265f3-4c62-4692-84ef-80747386eddd",
   "metadata": {},
   "source": [
    "### Setup LLM using Replicate\n",
    "\n",
    "Make sure you have REPLICATE_API_TOKEN specified!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b56a2ff-f9a4-4ba1-aee4-58f9f71637d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# os.environ[\"REPLICATE_API_TOKEN\"] = \"<YOUR_API_KEY>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60247c74-046b-4f6a-8522-bb4faf5a7a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.replicate import Replicate\n",
    "\n",
    "llm_replicate = Replicate(model=\"meta/meta-llama-3-70b-instruct\")\n",
    "# llm_replicate = Replicate(model=\"meta/meta-llama-3-8b-instruct\")"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b648635a-2672-407f-bae6-01660e5426d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id 32a7bb50-6a25-4295-971c-2de6f1588e0d\n",
      ".Started parsing the file under job_id b8cc075e-b6d5-4ded-b060-f72e9393b391\n",
      "..Started parsing the file under job_id 42fc41a4-68b6-49ee-8647-781b5cdb8893\n",
      "..."
     ]
    }
   ],
   "source": [
    "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",
    "\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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm just an AI, I don't have personal preferences or opinions, nor can I listen to music. I exist solely to provide information and assist with tasks, so I don't have the capacity to enjoy or compare different artists' music. Both Drake and Kendrick Lamar are highly acclaimed rappers, and it's subjective which one you might prefer based on your individual tastes in music.\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As a hypothetical Drake fan, I'd say that there are several reasons why I might prefer his music over Kendrick's. Here are a few possible reasons:\n",
      "\n",
      "1. **Lyrical storytelling**: Drake is known for his vivid storytelling on tracks like \"Marvins Room\" and \"Take Care.\" He has a way of painting pictures with his words, making listeners feel like they're right there with him, experiencing the highs and lows he's singing about. Kendrick, while also an incredible storyteller, might not have the same level of lyrical detail that Drake does.\n",
      "2. **Melodic flow**: Drake's melodic flow is infectious! He has a way of crafting hooks and choruses that get stuck in your head, making it hard to stop listening. Kendrick's flows are often more complex and intricate, but Drake's simplicity can be just as effective in getting the job done.\n",
      "3. **Vulnerability**: Drake isn't afraid to show his vulnerable side on tracks like \"Hold On\" and \"I'm Upset.\" He wears his heart on his sleeve, sharing personal struggles and emotions with listeners. This vulnerability makes him relatable and easier to connect with on a deeper level.\n",
      "4. **Production**: Drake has had the privilege of working with some incredible producers (like Noah \"40\" Shebib and Boi-1da) who bring out the best in him. The way he incorporates these sounds into his songs is often seamless, creating a unique blend of hip-hop and R&B that's hard to resist.\n",
      "5. **Cultural relevance**: As someone who grew up in Toronto, Drake has a deep understanding of the Canadian experience and the struggles that come with it. He often references his hometown and the people he grew up around, giving his music a distinctly Canadian flavor. This cultural relevance makes his music feel more authentic and connected to the world we live in.\n",
      "6. **Commercial appeal**: Let's face it – Drake has a knack for creating hits! His songs are often catchy, radio-friendly, and designed to get stuck in your head. While Kendrick might not have the same level of commercial success, Drake's ability to craft songs that resonate with a wider audience is undeniable.\n",
      "\n",
      "Of course, this is all just hypothetical – as a fan, I can appreciate both artists for their unique strengths and styles! What do you think?"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: \"Listen up, y'all, I got a message to share\n",
      "Been through the struggles, but my spirit's still fair\n",
      "From Compton streets to the top of the game\n",
      "I'm the real Hov, ain't nobody gonna claim my fame\"\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "According to the provided context, \"Family Matters\" is a seven-and-a-half-minute diss track by Drake in response to Kendrick Lamar's disses against him. The song has three different beats and features several shots at Kendrick, as well as other members of Drake's entourage, including A$AP Rocky and The Weeknd. In the song, Drake raps about his personal life, including his relationships with Rihanna and Whitney Alford, and even makes allegations about Kendrick's domestic life.\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**Repeat**\n",
      "\n",
      "The article does not provide a clear verdict on who \"won\" the beef, nor does it suggest that the conflict has been definitively resolved. Instead, it presents the situation as ongoing and multifaceted, with both artists continuing to engage in a game of verbal sparring and lyrical one-upmanship.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68918eb6-f1e6-460c-b1d5-fb49c3fed4b8",
   "metadata": {},
   "source": [
    "## 3. Advanced RAG (Routing, Sub-Questions)"
   ]
  },
  {
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mSelecting query engine 0: The song 'Meet the Grahams' might contain specific facts or information about the band, making it useful for searching for those specific details..\n",
      "\u001b[0m"
     ]
    }
   ],
   "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\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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"Meet the Grahams\" artwork is a crucial part of a larger strategy by Kendrick Lamar to address Drake's family matters in a diss track. The artwork shows a pair of Maybach gloves, a shirt, receipts, and prescription bottles, including one for Ozempic prescribed to Drake. This song is significant because it serves as the full picture that Kendrick teased earlier on \"6.16 in LA\" and addresses all members of Drake's family, including his son Adonis, mother Sandi, father Dennis, and an alleged 11-year-old daughter. The song takes it to the point of no return, with Kendrick musing that he wishes Dennis Graham wore a condom the night Drake was conceived and telling both Drake's parents that they raised a man whose house is due to be raided any day now on Harvey Weinstein-level allegations.\n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af04881-6a64-475c-8c1e-e5bd3e8006ee",
   "metadata": {},
   "source": [
    "### Break Complex Questions down into Sub-Questions\n",
    "\n",
    "Our Sub-Question Query Engine breaks complex questions down into sub-questions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9430df-ece7-493c-be6b-d5b371c7b8b0",
   "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": "f16fa210-7e15-474a-9bd2-1ac681290006",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
    "\n",
    "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_summary\",\n",
    "        description=\"Useful for searching over Kendrick's life.\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e05b95-95cb-45f3-8fbe-d0b54a9ba988",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 1 sub questions.\n",
      "\u001b[1;3;38;2;237;90;200m[drake_search] Q: What are the albums released by Drake\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200m[drake_search] A: Based on the provided context information, the albums released by Drake are:\n",
      "\n",
      "1. Take Care (album)\n",
      "2. Nothing Was the Same\n",
      "3. If You're Reading This It's Too Late (rumored to be a mixtape or album)\n",
      "4. Certified Lover Boy\n",
      "5. Honestly, Nevermind\n",
      "\u001b[0mBased on the provided context information, the albums released by Drake are:\n",
      "\n",
      "1. Take Care (album)\n",
      "2. Nothing Was the Same\n",
      "3. If You're Reading This It's Too Late (rumored to be a mixtape or album)\n",
      "4. Certified Lover Boy\n",
      "5. Honestly, Nevermind\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.query_engine import SubQuestionQueryEngine\n",
    "\n",
    "query_engine = SubQuestionQueryEngine.from_defaults(\n",
    "    [drake_tool, kendrick_tool],\n",
    "    llm=llm_replicate,  # llama3-70b\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"Which albums did Drake release in his career?\")\n",
    "\n",
    "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-05-10 23:40:37--  https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip\n",
      "Resolving www.sqlitetutorial.net (www.sqlitetutorial.net)... 2606:4700:3037::6815:1e8d, 2606:4700:3037::ac43:acfa, 104.21.30.141, ...\n",
      "Connecting to www.sqlitetutorial.net (www.sqlitetutorial.net)|2606:4700:3037::6815:1e8d|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 305596 (298K) [application/zip]\n",
      "Saving to: ‘./data/chinook.zip’\n",
      "\n",
      "./data/chinook.zip  100%[===================>] 298.43K  --.-KB/s    in 0.02s   \n",
      "\n",
      "2024-05-10 23:40:37 (13.9 MB/s) - ‘./data/chinook.zip’ saved [305596/305596]\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Archive:  ./data/chinook.zip\n",
      "  inflating: chinook.db              \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "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_replicate,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a65ecd70-09c4-4872-b712-3a8235d03db2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here are 10 album titles with their corresponding artists:\n",
      "\n",
      "1. \"For Those About To Rock We Salute You\" by Artist 1\n",
      "2. \"Balls to the Wall\" by Artist 2\n",
      "3. \"Restless and Wild\" by Artist 2\n",
      "4. \"Let There Be Rock\" by Artist 1\n",
      "5. \"Big Ones\" by Artist 3\n",
      "6. \"Jagged Little Pill\" by Artist 4\n",
      "7. \"Facelift\" by Artist 5\n",
      "8. \"Warner 25 Anos\" by Artist 6\n",
      "9. \"Plays Metallica By Four Cellos\" by Artist 7\n",
      "10. \"Audioslave\" by Artist 8\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here are 5 artists: AC/DC, Accept, Aerosmith, Alanis Morissette, and Alice In Chains.\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here are three tracks from the legendary Australian rock band AC/DC: \"For Those About To Rock (We Salute You)\", \"Put The Finger On You\", and \"Let's Get It Up\".\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT tracks.Name FROM tracks JOIN albums ON tracks.AlbumId = albums.AlbumId JOIN artists ON albums.ArtistId = artists.ArtistId WHERE artists.Name = 'AC/DC' LIMIT 3;\n"
     ]
    }
   ],
   "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 / Ollama, 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.ollama import Ollama\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 = Ollama(model=\"llama3\")\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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='Tropical Bites' city='Miami' cuisine='Caribbean'\n"
     ]
    }
   ],
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Kendrick Lamar's contributions to the beef!\n",
      "\n",
      "According to the article, Kendrick released several diss tracks in response to Drake's initial shots. One notable track is \"Not Like Us\", which directly addresses Drake and his perceived shortcomings.\n",
      "\n",
      "However, the article highlights that Kendrick's most significant response was his album \"Mr. Morale & The Big Steppers\", which features several tracks that can be seen as indirect disses towards Drake.\n",
      "\n",
      "The article also mentions that Kendrick's family has been a target of Drake's attacks, with Drake referencing Kendrick's estranged relationship with his partner Whitney and their two kids (one of whom is allegedly fathered by Dave Free).\n",
      "\n",
      "It's worth noting that Kendrick didn't directly respond to Drake's THP6 track. Instead, he focused on his own music and let the lyrics speak for themselves.\n",
      "\n",
      "Overall, Kendrick's approach was more subtle yet still packed a punch, showcasing his storytelling ability and lyrical prowess.\n",
      "\n",
      "Would you like me to elaborate on any specific tracks or moments from the beef?\n"
     ]
    }
   ],
   "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_replicate,\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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: add\n",
      "Action Input: {'a': 121, 'b': 2}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 123\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I have the result of the addition, now I need to multiply it by 5.\n",
      "Action: multiply\n",
      "Action Input: {'a': 123, 'b': 5}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 615\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: 615\n",
      "\u001b[0m615\n"
     ]
    }
   ],
   "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": "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,  ## TODO: define query tools\n",
    "    llm=llm_replicate,\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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: kendrick_search\n",
      "Action Input: {'input': \"Kendrick Lamar's childhood\"}\n",
      "\u001b[0m\u001b[1;3;34mObservation: Kendrick Lamar was born on June 17, 1987, in Compton, California. He is the first child of Kenneth \"Kenny\" Duckworth, a former gang hustler who previously worked at KFC, and Paula Oliver, a hairdresser who previously worked at McDonald's. Both of his parents are African Americans from the South Side of Chicago, and they relocated to Compton in 1984 due to his father's affiliation with the Gangster Disciples. Lamar was named after singer-songwriter Eddie Kendricks of the Temptations. He was an only child until the age of seven and was described as a loner by his mother.\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I have information about Kendrick's childhood, but I need to know more about Drake's upbringing to answer the question.\n",
      "Action: drake_search\n",
      "Action Input: {'input': \"Drake's childhood\"}\n",
      "\u001b[0m\u001b[1;3;34mObservation: Drake was raised in two neighborhoods. He lived on Weston Road in Toronto's working-class west end until grade six and attended Weston Memorial Junior Public School until grade four. He moved to one of the city's affluent neighbourhoods, Forest Hill, in 2000. Drake appeared in a comedic sketch which aired during the 1997 NHL Awards, featuring Martin Brodeur and Ron Hextall. At age 10, he attended Forest Hill Collegiate Institute for high school.\n",
      "\u001b[0m\u001b[1;3;34mObservation: Error: Could not parse output. Please follow the thought-action-input format. Try again.\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I apologize for the mistake. I need to use a tool to help me answer the question.\n",
      "Action: drake_search\n",
      "Action Input: {'input': \"Drake's childhood\"}\n",
      "\u001b[0m\u001b[1;3;34mObservation: Drake was raised in two neighborhoods. He lived on Weston Road in Toronto's working-class west end until grade six and attended Weston Memorial Junior Public School until grade four. He played minor hockey with the Weston Red Wings, reaching the Upper Canada College hockey camp before leaving due to a vicious cross-check to his neck during a game. At age 10, Drake appeared in a comedic sketch which aired during the 1997 NHL Awards.\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I have information about both Kendrick and Drake's childhood, so I can answer the question without using any more tools.\n",
      "Answer: Kendrick Lamar grew up in Compton, California, as the child of a former gang hustler and a hairdresser, while Drake was raised in two neighborhoods in Toronto, Ontario, Canada, and had a brief experience in minor hockey before pursuing a career in entertainment.\n",
      "\u001b[0mKendrick Lamar grew up in Compton, California, as the child of a former gang hustler and a hairdresser, while Drake was raised in two neighborhoods in Toronto, Ontario, Canada, and had a brief experience in minor hockey before pursuing a career in entertainment.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"Tell me about how Kendrick and Drake grew up\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v3",
   "language": "python",
   "name": "llama_index_v3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}