{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-recipes/blob/main/recipes/quickstart/agents/dlai/Building_Agentic_RAG_with_Llamaindex_L1_Router_Engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook ports the DeepLearning.AI short course [Building Agentic RAG with Llamaindex Lesson 1 Router Engine](https://learn.deeplearning.ai/courses/building-agentic-rag-with-llamaindex/lesson/2/router-query-engine) to using Llama 3. \n",
    "\n",
    "You should take the course before or after going through this notebook to have a deeper understanding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eiJsOa29ej7G",
    "outputId": "094c60c2-d782-4baf-bfc3-913f53ac1ff3"
   },
   "outputs": [],
   "source": [
    "!pip install llama-index\n",
    "!pip install llama-index-embeddings-huggingface\n",
    "!pip install llama-index-llms-groq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Wv_-jm5XGoWo",
    "outputId": "d0dd37f4-ca33-4704-d221-6ea80ee09eb5"
   },
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ['GROQ_API_KEY'] = 'your_groq_api_key' # get a free key at https://console.groq.com/keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aUi629qVUsG5",
    "outputId": "015a6a1a-a9e1-4c4f-dcaf-7d341b1a0b0e"
   },
   "outputs": [],
   "source": [
    "!wget \"https://openreview.net/pdf?id=VtmBAGCN7o\" -O metagpt.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "auZQalH5J7CU"
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(input_files=[\"metagpt.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GFfUjJypJ7Eq"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "splitter = SentenceSplitter(chunk_size=1024)\n",
    "nodes = splitter.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bguUp2D5LhST"
   },
   "outputs": [],
   "source": [
    "from llama_index.llms.groq import Groq\n",
    "\n",
    "from llama_index.core import Settings, VectorStoreIndex\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "llm = Groq(model=\"llama3-8b-8192\") #, api_key=GROQ_API_TOKEN)\n",
    "Settings.llm = llm\n",
    "#llm.complete(\"Who wrote the book godfather\").text\n",
    "\n",
    "Settings.embed_model = HuggingFaceEmbedding(\n",
    "    model_name=\"BAAI/bge-small-en-v1.5\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pkmnSHsMJ7Hg"
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SummaryIndex, VectorStoreIndex\n",
    "\n",
    "summary_index = SummaryIndex(nodes)\n",
    "vector_index = VectorStoreIndex(nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BnMq-qO9ezcE"
   },
   "outputs": [],
   "source": [
    "summary_query_engine = summary_index.as_query_engine(\n",
    "    response_mode=\"tree_summarize\",\n",
    "    use_async=True,\n",
    ")\n",
    "vector_query_engine = vector_index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vnK7YXlVLojh"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "summary_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=summary_query_engine,\n",
    "    description=(\n",
    "        \"Useful for summarization questions related to MetaGPT\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "vector_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=vector_query_engine,\n",
    "    description=(\n",
    "        \"Useful for retrieving specific context from the MetaGPT paper.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EvLJilU4LomU"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine.router_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=[\n",
    "        summary_tool,\n",
    "        vector_tool,\n",
    "    ],\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O_HfeD6TMCJf"
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KAjQW6M1LopK",
    "outputId": "fd85f083-599f-40b0-8fed-c3c96b09c5c9"
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the summary of the document?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ySNeKv3PLosR",
    "outputId": "9568d416-9ec1-4f50-dbbf-61a208205113"
   },
   "outputs": [],
   "source": [
    "print(len(response.source_nodes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KOP52yxpLou-",
    "outputId": "168a6ff3-2a2a-4588-dc2f-978ae8c0bbc1"
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\n",
    "    \"How do agents share information with other agents? This is not a summarization question.\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "luGziGzfMHbD"
   },
   "outputs": [],
   "source": [
    "def get_router_query_engine(file_path: str):\n",
    "    \"\"\"Get router query engine.\"\"\"\n",
    "\n",
    "    documents = SimpleDirectoryReader(input_files=[file_path]).load_data()\n",
    "\n",
    "    splitter = SentenceSplitter(chunk_size=1024)\n",
    "    nodes = splitter.get_nodes_from_documents(documents)\n",
    "\n",
    "    summary_index = SummaryIndex(nodes)\n",
    "    vector_index = VectorStoreIndex(nodes)\n",
    "\n",
    "    summary_query_engine = summary_index.as_query_engine(\n",
    "        response_mode=\"tree_summarize\",\n",
    "        use_async=True,\n",
    "    )\n",
    "    vector_query_engine = vector_index.as_query_engine()\n",
    "\n",
    "    summary_tool = QueryEngineTool.from_defaults(\n",
    "        query_engine=summary_query_engine,\n",
    "        description=(\n",
    "            \"Useful for summarization questions related to MetaGPT\"\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    vector_tool = QueryEngineTool.from_defaults(\n",
    "        query_engine=vector_query_engine,\n",
    "        description=(\n",
    "            \"Useful for retrieving specific context from the MetaGPT paper.\"\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    query_engine = RouterQueryEngine(\n",
    "        selector=LLMSingleSelector.from_defaults(),\n",
    "        query_engine_tools=[\n",
    "            summary_tool,\n",
    "            vector_tool,\n",
    "        ],\n",
    "        verbose=True\n",
    "    )\n",
    "    return query_engine\n",
    "\n",
    "query_engine = get_router_query_engine(\"metagpt.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SHqbTBxQNryl",
    "outputId": "ca2f16e9-2274-42e6-fdd5-3218d8f0c3f8"
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\"Tell me about the ablation study results?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5r1MHbLOPT8Y"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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": 4
}
