LlamaIndex_V0_Agentic_RAG_multi_document_agents.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "67c8afd7",
      "metadata": {
        "id": "67c8afd7"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/multi_document_agents.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "43497beb-817d-4366-9156-f4d7f0d44942",
      "metadata": {
        "id": "43497beb-817d-4366-9156-f4d7f0d44942"
      },
      "source": [
        "# Multi-Document Agents\n",
        "\n",
        "In this guide, you learn towards setting up an agent that can effectively answer different types of questions over a larger set of documents.\n",
        "\n",
        "These questions include the following\n",
        "\n",
        "- QA over a specific doc\n",
        "- QA comparing different docs\n",
        "- Summaries over a specific doc\n",
        "- Comparing summaries between different docs\n",
        "\n",
        "We do this with the following architecture:\n",
        "\n",
        "- setup a \"document agent\" over each Document: each doc agent can do QA/summarization within its doc\n",
        "- setup a top-level agent over this set of document agents. Do tool retrieval and then do CoT over the set of tools to answer a question."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9be00aba-b6c5-4940-9825-81c5d2cd2f0b",
      "metadata": {
        "id": "9be00aba-b6c5-4940-9825-81c5d2cd2f0b"
      },
      "source": [
        "## Setup and Download Data\n",
        "\n",
        "In this section, we'll define imports and then download Wikipedia articles about different cities. Each article is stored separately.\n",
        "\n",
        "We load in 18 cities - this is not quite at the level of \"hundreds\" of documents but its still large enough to warrant some top-level document retrieval!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5d81f93c",
      "metadata": {
        "id": "5d81f93c"
      },
      "source": [
        "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f0fb1340",
      "metadata": {
        "id": "f0fb1340",
        "outputId": "e7792337-06a7-4de4-914d-52a70c0b5351",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting llama-index\n",
            "  Downloading llama_index-0.9.39-py3-none-any.whl (15.9 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m15.9/15.9 MB\u001b[0m \u001b[31m30.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: SQLAlchemy[asyncio]>=1.4.49 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2.0.24)\n",
            "Requirement already satisfied: aiohttp<4.0.0,>=3.8.6 in /usr/local/lib/python3.10/dist-packages (from llama-index) (3.9.1)\n",
            "Collecting dataclasses-json (from llama-index)\n",
            "  Downloading dataclasses_json-0.6.3-py3-none-any.whl (28 kB)\n",
            "Collecting deprecated>=1.2.9.3 (from llama-index)\n",
            "  Downloading Deprecated-1.2.14-py2.py3-none-any.whl (9.6 kB)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2023.6.0)\n",
            "Collecting httpx (from llama-index)\n",
            "  Downloading httpx-0.26.0-py3-none-any.whl (75 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.9/75.9 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: nest-asyncio<2.0.0,>=1.5.8 in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.6.0)\n",
            "Requirement already satisfied: networkx>=3.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (3.2.1)\n",
            "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /usr/local/lib/python3.10/dist-packages (from llama-index) (3.8.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.23.5)\n",
            "Collecting openai>=1.1.0 (from llama-index)\n",
            "  Downloading openai-1.10.0-py3-none-any.whl (225 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m225.1/225.1 kB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.5.3)\n",
            "Requirement already satisfied: requests>=2.31.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2.31.0)\n",
            "Requirement already satisfied: tenacity<9.0.0,>=8.2.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (8.2.3)\n",
            "Collecting tiktoken>=0.3.3 (from llama-index)\n",
            "  Downloading tiktoken-0.5.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m40.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (4.5.0)\n",
            "Collecting typing-inspect>=0.8.0 (from llama-index)\n",
            "  Downloading typing_inspect-0.9.0-py3-none-any.whl (8.8 kB)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (23.2.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (1.9.4)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (1.4.1)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (1.3.1)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.6->llama-index) (4.0.3)\n",
            "Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.10/dist-packages (from deprecated>=1.2.9.3->llama-index) (1.14.1)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (8.1.7)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (1.3.2)\n",
            "Requirement already satisfied: regex>=2021.8.3 in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (2023.6.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (4.66.1)\n",
            "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.10/dist-packages (from openai>=1.1.0->llama-index) (3.7.1)\n",
            "Requirement already satisfied: distro<2,>=1.7.0 in /usr/lib/python3/dist-packages (from openai>=1.1.0->llama-index) (1.7.0)\n",
            "Requirement already satisfied: pydantic<3,>=1.9.0 in /usr/local/lib/python3.10/dist-packages (from openai>=1.1.0->llama-index) (1.10.14)\n",
            "Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from openai>=1.1.0->llama-index) (1.3.0)\n",
            "Collecting typing-extensions>=4.5.0 (from llama-index)\n",
            "  Downloading typing_extensions-4.9.0-py3-none-any.whl (32 kB)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx->llama-index) (2023.11.17)\n",
            "Collecting httpcore==1.* (from httpx->llama-index)\n",
            "  Downloading httpcore-1.0.2-py3-none-any.whl (76 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.9/76.9 kB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: idna in /usr/local/lib/python3.10/dist-packages (from httpx->llama-index) (3.6)\n",
            "Collecting h11<0.15,>=0.13 (from httpcore==1.*->httpx->llama-index)\n",
            "  Downloading h11-0.14.0-py3-none-any.whl (58 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.31.0->llama-index) (3.3.2)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.31.0->llama-index) (2.0.7)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy[asyncio]>=1.4.49->llama-index) (3.0.3)\n",
            "Collecting mypy-extensions>=0.3.0 (from typing-inspect>=0.8.0->llama-index)\n",
            "  Downloading mypy_extensions-1.0.0-py3-none-any.whl (4.7 kB)\n",
            "Collecting marshmallow<4.0.0,>=3.18.0 (from dataclasses-json->llama-index)\n",
            "  Downloading marshmallow-3.20.2-py3-none-any.whl (49 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas->llama-index) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->llama-index) (2023.3.post1)\n",
            "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai>=1.1.0->llama-index) (1.2.0)\n",
            "Requirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.10/dist-packages (from marshmallow<4.0.0,>=3.18.0->dataclasses-json->llama-index) (23.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.1->pandas->llama-index) (1.16.0)\n",
            "Installing collected packages: typing-extensions, mypy-extensions, marshmallow, h11, deprecated, typing-inspect, tiktoken, httpcore, httpx, dataclasses-json, openai, llama-index\n",
            "  Attempting uninstall: typing-extensions\n",
            "    Found existing installation: typing_extensions 4.5.0\n",
            "    Uninstalling typing_extensions-4.5.0:\n",
            "      Successfully uninstalled typing_extensions-4.5.0\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "llmx 0.0.15a0 requires cohere, which is not installed.\n",
            "tensorflow-probability 0.22.0 requires typing-extensions<4.6.0, but you have typing-extensions 4.9.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed dataclasses-json-0.6.3 deprecated-1.2.14 h11-0.14.0 httpcore-1.0.2 httpx-0.26.0 llama-index-0.9.39 marshmallow-3.20.2 mypy-extensions-1.0.0 openai-1.10.0 tiktoken-0.5.2 typing-extensions-4.9.0 typing-inspect-0.9.0\n"
          ]
        }
      ],
      "source": [
        "!pip install llama-index"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install openai"
      ],
      "metadata": {
        "id": "JC4Kpp227wCn",
        "outputId": "2ed26cc5-6bc0-4404-d05e-04f493dd586f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "id": "JC4Kpp227wCn",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: openai in /usr/local/lib/python3.10/dist-packages (1.10.0)\n",
            "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.10/dist-packages (from openai) (3.7.1)\n",
            "Requirement already satisfied: distro<2,>=1.7.0 in /usr/lib/python3/dist-packages (from openai) (1.7.0)\n",
            "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from openai) (0.26.0)\n",
            "Requirement already satisfied: pydantic<3,>=1.9.0 in /usr/local/lib/python3.10/dist-packages (from openai) (1.10.14)\n",
            "Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from openai) (1.3.0)\n",
            "Requirement already satisfied: tqdm>4 in /usr/local/lib/python3.10/dist-packages (from openai) (4.66.1)\n",
            "Requirement already satisfied: typing-extensions<5,>=4.7 in /usr/local/lib/python3.10/dist-packages (from openai) (4.9.0)\n",
            "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai) (3.6)\n",
            "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai) (1.2.0)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx<1,>=0.23.0->openai) (2023.11.17)\n",
            "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx<1,>=0.23.0->openai) (1.0.2)\n",
            "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.10/dist-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.14.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import openai\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"Your API Key goes here\""
      ],
      "metadata": {
        "id": "dTtgdPb47z8h"
      },
      "id": "dTtgdPb47z8h",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e41e9905-77a9-44c5-88ac-c7a4d08a4612",
      "metadata": {
        "id": "e41e9905-77a9-44c5-88ac-c7a4d08a4612"
      },
      "outputs": [],
      "source": [
        "from llama_index import (\n",
        "    VectorStoreIndex,\n",
        "    SummaryIndex,\n",
        "    SimpleKeywordTableIndex,\n",
        "    SimpleDirectoryReader,\n",
        "    ServiceContext,\n",
        ")\n",
        "from llama_index.schema import IndexNode\n",
        "from llama_index.tools import QueryEngineTool, ToolMetadata\n",
        "from llama_index.llms import OpenAI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e4343cf7-eec9-4a67-b5be-c72dbe3280a7",
      "metadata": {
        "id": "e4343cf7-eec9-4a67-b5be-c72dbe3280a7"
      },
      "outputs": [],
      "source": [
        "wiki_titles = [\n",
        "    \"Toronto\",\n",
        "    \"Seattle\",\n",
        "    \"Chicago\",\n",
        "    \"Boston\",\n",
        "    \"Houston\",\n",
        "    \"Tokyo\",\n",
        "    \"Berlin\",\n",
        "    \"Lisbon\",\n",
        "    \"Paris\",\n",
        "    \"London\",\n",
        "    \"Atlanta\",\n",
        "    \"Munich\",\n",
        "    \"Shanghai\",\n",
        "    \"Beijing\",\n",
        "    \"Copenhagen\",\n",
        "    \"Moscow\",\n",
        "    \"Cairo\",\n",
        "    \"Karachi\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6d261882-6793-4eca-ad93-d94d2061e388",
      "metadata": {
        "id": "6d261882-6793-4eca-ad93-d94d2061e388"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "\n",
        "import requests\n",
        "\n",
        "for title in wiki_titles:\n",
        "    response = requests.get(\n",
        "        \"https://en.wikipedia.org/w/api.php\",\n",
        "        params={\n",
        "            \"action\": \"query\",\n",
        "            \"format\": \"json\",\n",
        "            \"titles\": title,\n",
        "            \"prop\": \"extracts\",\n",
        "            # 'exintro': True,\n",
        "            \"explaintext\": True,\n",
        "        },\n",
        "    ).json()\n",
        "    page = next(iter(response[\"query\"][\"pages\"].values()))\n",
        "    wiki_text = page[\"extract\"]\n",
        "\n",
        "    data_path = Path(\"data\")\n",
        "    if not data_path.exists():\n",
        "        Path.mkdir(data_path)\n",
        "\n",
        "    with open(data_path / f\"{title}.txt\", \"w\") as fp:\n",
        "        fp.write(wiki_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5bf0c13b-0d77-43e8-8c1c-84258299a494",
      "metadata": {
        "id": "5bf0c13b-0d77-43e8-8c1c-84258299a494"
      },
      "outputs": [],
      "source": [
        "# Load all wiki documents\n",
        "city_docs = {}\n",
        "for wiki_title in wiki_titles:\n",
        "    city_docs[wiki_title] = SimpleDirectoryReader(\n",
        "        input_files=[f\"data/{wiki_title}.txt\"]\n",
        "    ).load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6189aaf4-2eb7-40bc-9e83-79ce4f221b4b",
      "metadata": {
        "id": "6189aaf4-2eb7-40bc-9e83-79ce4f221b4b"
      },
      "source": [
        "Define LLM + Service Context + Callback Manager"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd6e5e48-91b9-4701-a85d-d98c92323350",
      "metadata": {
        "id": "dd6e5e48-91b9-4701-a85d-d98c92323350"
      },
      "outputs": [],
      "source": [
        "llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
        "service_context = ServiceContext.from_defaults(llm=llm)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4eeef31a-fc25-4367-a5ba-945f81d04cf9",
      "metadata": {
        "id": "4eeef31a-fc25-4367-a5ba-945f81d04cf9"
      },
      "source": [
        "## Building Multi-Document Agents\n",
        "\n",
        "In this section we show you how to construct the multi-document agent. We first build a document agent for each document, and then define the top-level parent agent with an object index."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "976cd798-2e8d-474c-922a-51b12c5c6f36",
      "metadata": {
        "id": "976cd798-2e8d-474c-922a-51b12c5c6f36"
      },
      "source": [
        "### Build Document Agent for each Document\n",
        "\n",
        "In this section we define \"document agents\" for each document.\n",
        "\n",
        "We define both a vector index (for semantic search) and summary index (for summarization) for each document. The two query engines are then converted into tools that are passed to an OpenAI function calling agent.\n",
        "\n",
        "This document agent can dynamically choose to perform semantic search or summarization within a given document.\n",
        "\n",
        "We create a separate document agent for each city."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eacdf3a7-cfe3-4c2b-9037-b28a065ed148",
      "metadata": {
        "id": "eacdf3a7-cfe3-4c2b-9037-b28a065ed148"
      },
      "outputs": [],
      "source": [
        "from llama_index.agent import OpenAIAgent\n",
        "from llama_index import load_index_from_storage, StorageContext\n",
        "from llama_index.node_parser import SentenceSplitter\n",
        "import os\n",
        "\n",
        "node_parser = SentenceSplitter()\n",
        "\n",
        "# Build agents dictionary\n",
        "agents = {}\n",
        "query_engines = {}\n",
        "\n",
        "# this is for the baseline\n",
        "all_nodes = []\n",
        "\n",
        "for idx, wiki_title in enumerate(wiki_titles):\n",
        "    nodes = node_parser.get_nodes_from_documents(city_docs[wiki_title])\n",
        "    all_nodes.extend(nodes)\n",
        "\n",
        "    if not os.path.exists(f\"./data/{wiki_title}\"):\n",
        "        # build vector index\n",
        "        vector_index = VectorStoreIndex(nodes, service_context=service_context)\n",
        "        vector_index.storage_context.persist(\n",
        "            persist_dir=f\"./data/{wiki_title}\"\n",
        "        )\n",
        "    else:\n",
        "        vector_index = load_index_from_storage(\n",
        "            StorageContext.from_defaults(persist_dir=f\"./data/{wiki_title}\"),\n",
        "            service_context=service_context,\n",
        "        )\n",
        "\n",
        "    # build summary index\n",
        "    summary_index = SummaryIndex(nodes, service_context=service_context)\n",
        "    # define query engines\n",
        "    vector_query_engine = vector_index.as_query_engine()\n",
        "    summary_query_engine = summary_index.as_query_engine()\n",
        "\n",
        "    # define tools\n",
        "    query_engine_tools = [\n",
        "        QueryEngineTool(\n",
        "            query_engine=vector_query_engine,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"vector_tool\",\n",
        "                description=(\n",
        "                    \"Useful for questions related to specific aspects of\"\n",
        "                    f\" {wiki_title} (e.g. the history, arts and culture,\"\n",
        "                    \" sports, demographics, or more).\"\n",
        "                ),\n",
        "            ),\n",
        "        ),\n",
        "        QueryEngineTool(\n",
        "            query_engine=summary_query_engine,\n",
        "            metadata=ToolMetadata(\n",
        "                name=\"summary_tool\",\n",
        "                description=(\n",
        "                    \"Useful for any requests that require a holistic summary\"\n",
        "                    f\" of EVERYTHING about {wiki_title}. For questions about\"\n",
        "                    \" more specific sections, please use the vector_tool.\"\n",
        "                ),\n",
        "            ),\n",
        "        ),\n",
        "    ]\n",
        "\n",
        "    # build agent\n",
        "    function_llm = OpenAI(model=\"gpt-4\")\n",
        "    agent = OpenAIAgent.from_tools(\n",
        "        query_engine_tools,\n",
        "        llm=function_llm,\n",
        "        verbose=True,\n",
        "        system_prompt=f\"\"\"\\\n",
        "You are a specialized agent designed to answer queries about {wiki_title}.\n",
        "You must ALWAYS use at least one of the tools provided when answering a question; do NOT rely on prior knowledge.\\\n",
        "\"\"\",\n",
        "    )\n",
        "\n",
        "    agents[wiki_title] = agent\n",
        "    query_engines[wiki_title] = vector_index.as_query_engine(\n",
        "        similarity_top_k=2\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "899ca55b-0c02-429b-a765-8e4f806d503f",
      "metadata": {
        "id": "899ca55b-0c02-429b-a765-8e4f806d503f"
      },
      "source": [
        "### Build Retriever-Enabled OpenAI Agent\n",
        "\n",
        "We build a top-level agent that can orchestrate across the different document agents to answer any user query.\n",
        "\n",
        "This agent takes in all document agents as tools. This specific agent `RetrieverOpenAIAgent` performs tool retrieval before tool use (unlike a default agent that tries to put all tools in the prompt).\n",
        "\n",
        "Here we use a top-k retriever, but we encourage you to customize the tool retriever method!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6884ff15-bf40-4bdd-a1e3-58cbd056a12a",
      "metadata": {
        "id": "6884ff15-bf40-4bdd-a1e3-58cbd056a12a"
      },
      "outputs": [],
      "source": [
        "# define tool for each document agent\n",
        "all_tools = []\n",
        "for wiki_title in wiki_titles:\n",
        "    wiki_summary = (\n",
        "        f\"This content contains Wikipedia articles about {wiki_title}. Use\"\n",
        "        f\" this tool if you want to answer any questions about {wiki_title}.\\n\"\n",
        "    )\n",
        "    doc_tool = QueryEngineTool(\n",
        "        query_engine=agents[wiki_title],\n",
        "        metadata=ToolMetadata(\n",
        "            name=f\"tool_{wiki_title}\",\n",
        "            description=wiki_summary,\n",
        "        ),\n",
        "    )\n",
        "    all_tools.append(doc_tool)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "da80793c-54b9-43e2-b34a-27412156268a",
      "metadata": {
        "id": "da80793c-54b9-43e2-b34a-27412156268a"
      },
      "outputs": [],
      "source": [
        "# define an \"object\" index and retriever over these tools\n",
        "from llama_index import VectorStoreIndex\n",
        "from llama_index.objects import ObjectIndex, SimpleToolNodeMapping\n",
        "\n",
        "tool_mapping = SimpleToolNodeMapping.from_objects(all_tools)\n",
        "obj_index = ObjectIndex.from_objects(\n",
        "    all_tools,\n",
        "    tool_mapping,\n",
        "    VectorStoreIndex,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fed38942-1e37-4c61-89fa-d2ef41151831",
      "metadata": {
        "id": "fed38942-1e37-4c61-89fa-d2ef41151831"
      },
      "outputs": [],
      "source": [
        "from llama_index.agent import FnRetrieverOpenAIAgent\n",
        "\n",
        "top_agent = FnRetrieverOpenAIAgent.from_retriever(\n",
        "    obj_index.as_retriever(similarity_top_k=3),\n",
        "    system_prompt=\"\"\" \\\n",
        "You are an agent designed to answer queries about a set of given cities.\n",
        "Please always use the tools provided to answer a question. Do not rely on prior knowledge.\\\n",
        "\n",
        "\"\"\",\n",
        "    verbose=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aa32b97c-6779-4b60-823d-6ca3be6f358a",
      "metadata": {
        "id": "aa32b97c-6779-4b60-823d-6ca3be6f358a"
      },
      "source": [
        "### Define Baseline Vector Store Index\n",
        "\n",
        "As a point of comparison, we define a \"naive\" RAG pipeline which dumps all docs into a single vector index collection.\n",
        "\n",
        "We set the top_k = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "78e7b48c-687a-4585-a730-e651212de3b5",
      "metadata": {
        "id": "78e7b48c-687a-4585-a730-e651212de3b5"
      },
      "outputs": [],
      "source": [
        "base_index = VectorStoreIndex(all_nodes)\n",
        "base_query_engine = base_index.as_query_engine(similarity_top_k=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8dedb927-a992-4f21-a0fb-4ce4361adcb3",
      "metadata": {
        "id": "8dedb927-a992-4f21-a0fb-4ce4361adcb3"
      },
      "source": [
        "## Running Example Queries\n",
        "\n",
        "Let's run some example queries, ranging from QA / summaries over a single document to QA / summarization over multiple documents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e743c62-7dd8-4ac9-85a5-f1cbc112a79c",
      "metadata": {
        "id": "8e743c62-7dd8-4ac9-85a5-f1cbc112a79c",
        "outputId": "7f17dfda-1544-4be7-8ab6-5a23094807f2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Boston with args: {\n",
            "  \"input\": \"arts and culture\"\n",
            "}\n",
            "Added user message to memory: arts and culture\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"arts and culture\"\n",
            "}\n",
            "Got output: Boston has a rich arts and culture scene. The city is home to several art museums and galleries, including the Museum of Fine Arts and the Isabella Stewart Gardner Museum. The Institute of Contemporary Art is also located in Boston's Seaport District. The city has a vibrant theater district, with several theaters such as the Cutler Majestic Theatre, Citi Performing Arts Center, the Colonial Theater, and the Orpheum Theatre. Boston is known for its annual events, including the Boston Early Music Festival, the Boston Arts Festival, and the Italian summer feasts in the North End. The city also has a strong music culture, with the Boston Symphony Orchestra being one of the \"Big Five\" American orchestras. Overall, Boston offers a diverse range of arts and cultural experiences for residents and visitors alike.\n",
            "========================\n",
            "\n",
            "Got output: Boston is renowned for its vibrant arts and culture scene. The city houses several prestigious art museums and galleries, such as the Museum of Fine Arts and the Isabella Stewart Gardner Museum. The Institute of Contemporary Art, located in Boston's Seaport District, is another notable cultural hub.\n",
            "\n",
            "Boston's theater district is bustling with activity, featuring several theaters like the Cutler Majestic Theatre, Citi Performing Arts Center, the Colonial Theater, and the Orpheum Theatre. These venues host a variety of performances, from Broadway shows to local productions.\n",
            "\n",
            "The city is also known for its annual events that celebrate arts and culture. These include the Boston Early Music Festival, the Boston Arts Festival, and the Italian summer feasts in the North End.\n",
            "\n",
            "Music is a significant part of Boston's cultural identity. The Boston Symphony Orchestra, recognized as one of the \"Big Five\" American orchestras, is a testament to the city's rich musical heritage.\n",
            "\n",
            "In summary, Boston offers a diverse range of arts and cultural experiences for both its residents and visitors.\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# should use Boston agent -> vector tool\n",
        "response = top_agent.query(\"Tell me about the arts and culture in Boston\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a4ce2a76-5779-4acf-9337-69109dae7fd6",
      "metadata": {
        "id": "a4ce2a76-5779-4acf-9337-69109dae7fd6",
        "outputId": "ccc9e0bf-8a4b-4f9e-d61d-3ef93509c94c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Boston has a vibrant arts and culture scene with several renowned art museums and galleries. The city is home to the Museum of Fine Arts and the Isabella Stewart Gardner Museum, both of which showcase a wide range of artistic works. The Institute of Contemporary Art, located in the Seaport District, is another popular destination for contemporary art enthusiasts.\n",
            "\n",
            "Boston's theater district is also a hub for performing arts. The district is home to several theaters, including the Cutler Majestic Theatre, Citi Performing Arts Center, Colonial Theater, and Orpheum Theatre. These venues host a variety of performances, including Broadway shows, musicals, and plays.\n",
            "\n",
            "The city is known for its annual events that celebrate arts and culture. The Boston Early Music Festival showcases early music performances and attracts musicians and enthusiasts from around the world. The Boston Arts Festival is a multi-day event that features local artists, musicians, and performers. Additionally, the North End neighborhood hosts Italian summer feasts that celebrate Italian culture through food, music, and parades.\n",
            "\n",
            "Music is an integral part of Boston's cultural identity. The Boston Symphony Orchestra, one of the \"Big Five\" American orchestras, performs at the renowned Symphony Hall. The city also has a thriving live music scene with numerous venues that host local and touring musicians.\n",
            "\n",
            "Overall, Boston offers a diverse range of arts and cultural experiences, making it a vibrant and dynamic city for art lovers and enthusiasts.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "af28b422-fb73-4b59-9e77-3ba3afa87795",
      "metadata": {
        "id": "af28b422-fb73-4b59-9e77-3ba3afa87795",
        "outputId": "2f598985-c1cc-4512-db44-4bd4f113387d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Boston has a rich arts and culture scene. The city is known for its literary culture, with famous writers like Ralph Waldo Emerson, Henry David Thoreau, and Nathaniel Hawthorne having roots in Boston. The Old Corner Bookstore is considered the \"cradle of American literature\" and the place where these writers met. The Boston Public Library, founded in 1852, was the first free library in the United States. Today, Boston's literary culture thrives with the presence of many universities and the annual Boston Book Festival.\n",
            "\n",
            "Music is also highly valued in Boston, with the Boston Symphony Orchestra being one of the \"Big Five\" American orchestras. Symphony Hall, home to the Boston Symphony Orchestra, is considered one of the top venues for classical music in the world. The city is also home to the Boston Pops Orchestra and the Boston Youth Symphony Orchestra, the largest youth orchestra in the nation. Other performing arts organizations in Boston include the Boston Ballet, Boston Lyric Opera Company, and the Handel and Haydn Society.\n",
            "\n",
            "Boston is home to several art museums and galleries, including the Museum of Fine Arts and the Isabella Stewart Gardner Museum. The city also has a vibrant contemporary art scene, with the Institute of Contemporary Art showcasing modern art in a contemporary building. The South End Art and Design District\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell me about the arts and culture in Boston\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "aa3d98ab-cb82-4473-ab2b-bc8a17e1b86a",
      "metadata": {
        "id": "aa3d98ab-cb82-4473-ab2b-bc8a17e1b86a",
        "outputId": "ae9497e6-544d-4ff2-d452-80ae7fe94c01",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Houston with args: {\n",
            "  \"input\": \"positive aspects\"\n",
            "}\n",
            "Added user message to memory: positive aspects\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"positive aspects\"\n",
            "}\n",
            "Got output: Houston has several positive aspects. It is recognized worldwide for its energy industry, particularly for oil and natural gas. The city is also known for its biomedical research and aeronautics. In addition, Houston is a growing hub for technology startup firms and has a fast-growing technology sector. The city's economy is diverse and includes major technology and software companies. Houston is also a global city and a top U.S. market for exports. The Houston area has a strong petrochemical industry and is home to the Port of Houston, which ranks first in international commerce in the United States. The city's gross domestic product (GDP) is one of the largest in the United States and larger than several countries. The University of Houston System also has a significant impact on the local economy, attracting new funds and generating jobs. Houston has been recognized for its business and career opportunities, reasonable living costs, and quality of life. It has also been ranked highly for technological innovation, employment creation, and shopping.\n",
            "========================\n",
            "\n",
            "Got output: Houston has numerous positive aspects that make it a vibrant and dynamic city. \n",
            "\n",
            "1. **Economy**: Houston is globally recognized for its energy industry, particularly oil and natural gas. It's also known for its biomedical research and aeronautics. The city is a growing hub for technology startups and has a fast-growing technology sector. \n",
            "\n",
            "2. **Diversity**: The city's economy is diverse, housing major technology and software companies. Houston is a global city and a top U.S. market for exports. \n",
            "\n",
            "3. **Petrochemical Industry**: The Houston area has a strong petrochemical industry and is home to the Port of Houston, which ranks first in international commerce in the United States. \n",
            "\n",
            "4. **GDP**: The city's gross domestic product (GDP) is one of the largest in the United States and larger than several countries. \n",
            "\n",
            "5. **Education**: The University of Houston System has a significant impact on the local economy, attracting new funds and generating jobs. \n",
            "\n",
            "6. **Quality of Life**: Houston has been recognized for its business and career opportunities, reasonable living costs, and quality of life. It has also been ranked highly for technological innovation, employment creation, and shopping.\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# should use Houston agent -> vector tool\n",
        "response = top_agent.query(\n",
        "    \"Give me a summary of all the positive aspects of Houston\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d476c54b-98af-4d2a-8f17-4baa37d0d360",
      "metadata": {
        "id": "d476c54b-98af-4d2a-8f17-4baa37d0d360",
        "outputId": "8ebe9dc5-2cd5-4c14-8b83-eea513cc7da6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Houston has numerous positive aspects that make it a vibrant and dynamic city. \n",
            "\n",
            "1. Economy: Houston is globally recognized for its energy industry, particularly oil and natural gas. It's also known for its biomedical research and aeronautics. The city is a growing hub for technology startups and has a fast-growing technology sector. \n",
            "\n",
            "2. Diversity: The city's economy is diverse, housing major technology and software companies. Houston is a global city and a top U.S. market for exports. \n",
            "\n",
            "3. Petrochemical Industry: The Houston area has a strong petrochemical industry and is home to the Port of Houston, which ranks first in international commerce in the United States. \n",
            "\n",
            "4. GDP: The city's gross domestic product (GDP) is one of the largest in the United States and larger than several countries. \n",
            "\n",
            "5. Education: The University of Houston System has a significant impact on the local economy, attracting new funds and generating jobs. \n",
            "\n",
            "6. Quality of Life: Houston has been recognized for its business and career opportunities, reasonable living costs, and quality of life. It has also been ranked highly for technological innovation, employment creation, and shopping.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c89b3581-ff70-4069-8cbb-f2791756128a",
      "metadata": {
        "id": "c89b3581-ff70-4069-8cbb-f2791756128a",
        "outputId": "1c094539-e0c7-45c2-dd97-449ae5235217",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Houston has several positive aspects. It has experienced significant growth and economic success, with a thriving energy industry, biomedical research, and aeronautics. The city is recognized as a global city and is a top market for exports in the United States. Houston is also known for its diverse culture and large international community. It has a vibrant arts and theater scene, with numerous performing arts organizations and museums. The city offers a variety of recreational opportunities, including parks, green spaces, and tourist attractions like Space Center Houston. Additionally, Houston has been recognized for its food and restaurant culture, and it has been ranked highly in various lists for its technological innovation, job creation, and quality of life.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Give me a summary of all the positive aspects of Houston\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c9d7dcec-5d1a-485a-b40f-aa9604fd92df",
      "metadata": {
        "id": "c9d7dcec-5d1a-485a-b40f-aa9604fd92df",
        "outputId": "85352c4c-0d0d-42c0-9138-60235a2ab2ad",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 469
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'== Economy ==\\n\\nHouston is recognized worldwide for its energy industry—particularly for oil and natural gas—as well as for biomedical research and aeronautics. Renewable energy sources—wind and solar—are also growing economic bases in the city, and the City Government purchases 90% of its annual 1 TWh power mostly from wind, and some from solar. The city has also been a growing hub for technology startup firms and is the fastest growing sector of the city\\'s economy. Major technology and software companies within Greater Houston include Crown Castle, KBR, FlightAware, Cybersoft, Houston Wire & Cable, and HostGator. Aylo, Go Daddy, and ByteDance have offices in the Houston area. On April 4, 2022, Hewlett Packard Enterprise relocated its global headquarters from California to the Greater Houston area. The Houston Ship Channel is also a large part of Houston\\'s economic base.\\nBecause of these strengths, Houston is designated as a global city by the Globalization and World Cities Study Group and Network and global management consulting firm A.T. Kearney. The Houston area is the top U.S. market for exports, surpassing New York City in 2013, according to data released by the U.S. Department of Commerce\\'s International Trade Administration. In 2012, the Houston–The Woodlands–Sugar Land area recorded $110.3 billion in merchandise exports. Petroleum products, chemicals, and oil and gas extraction equipment accounted for roughly two-thirds of the metropolitan area\\'s exports last year. The top three destinations for exports were Mexico, Canada, and Brazil.The Houston area is a leading center for building oilfield equipment. Much of its success as a petrochemical complex is due to its busy ship channel, the Port of Houston. In the United States, the port ranks first in international commerce and 16th among the largest ports in the world. Unlike most places, high oil and gasoline prices are beneficial for Houston\\'s economy, as many of its residents are employed in the energy industry. Houston is the beginning or end point of numerous oil, gas, and products pipelines.The Houston–The Woodlands–Sugar Land metro area\\'s gross domestic product (GDP) in 2022 was $633 billion, making it the seventh-largest of any metropolitan area in the United States and larger than Iran\\'s, Colombia\\'s, or the United Arab Emirates\\' GDP. Only 27 countries other than the United States have a gross domestic product exceeding Houston\\'s regional gross area product (GAP). In 2010, mining (which consists almost entirely of exploration and production of oil and gas in Houston) accounted for 26.3% of Houston\\'s GAP up sharply in response to high energy prices and a decreased worldwide surplus of oil production capacity, followed by engineering services, health services, and manufacturing.The University of Houston System\\'s annual impact on the Houston area\\'s economy equates to that of a major corporation: $1.1 billion in new funds attracted annually to the Houston area, $3.13 billion in total economic benefit, and 24,000 local jobs generated. This is in addition to the 12,500 new graduates the U.H. System produces every year who enter the workforce in Houston and throughout Texas. These degree-holders tend to stay in Houston. After five years, 80.5% of graduates are still living and working in the region.In 2019, the Houston metropolitan area ranked third in Texas within the category of \"Best Places for Business and Careers\" by Forbes magazine. Ninety-one foreign governments have established consular offices in Houston\\'s metropolitan area, the third-highest in the nation. Forty foreign governments maintain trade and commercial offices here with 23 active foreign chambers of commerce and trade associations. Twenty-five foreign banks representing 13 nations operate in Houston, providing financial assistance to the international community.In 2008, Houston received top ranking on Kiplinger\\'s Personal Finance \"Best Cities of 2008\" list, which ranks cities on their local economy, employment opportunities, reasonable living costs, and quality of life. The city ranked fourth for highest increase in the local technological innovation over the preceding 15 years, according to Forbes magazine. In the same year, the city ranked second on the annual Fortune 500 list of company headquarters, first for Forbes magazine\\'s \"Best Cities for College Graduates\", and first on their list of \"Best Cities to Buy a Home\". In 2010, the city was rated the best city for shopping, according to Forbes.In 2012, the city was ranked number one for paycheck worth by Forbes and in late May 2013, Houston was identified as America\\'s top city for employment creation.In 2013, Houston was identified as the number one U.S. city for job creation by the U.S.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 21
        }
      ],
      "source": [
        "# baseline: the response doesn't quite match the sources...\n",
        "response.source_nodes[1].get_content()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee6ef20c-3ccc-46c3-ad87-667138d78d5d",
      "metadata": {
        "id": "ee6ef20c-3ccc-46c3-ad87-667138d78d5d",
        "outputId": "13cb2fab-16bb-4d6d-a072-791aa50ec84d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Houston with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Added user message to memory: demographics\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Got output: Houston has a population of 2,304,580 according to the 2020 U.S. census. In 2017, the estimated population was 2,312,717, and in 2018 it was 2,325,502. The city has a diverse demographic makeup, with a significant number of undocumented immigrants residing in the Houston area, comprising nearly 9% of the city's metropolitan population in 2017. The age distribution in Houston shows a significant number of residents under 15 and between the ages of 20 to 34. The median age of the city is 33.4. The city has a mix of homeowners and renters, with an estimated 42.3% of Houstonians owning housing units. The median household income in 2019 was $52,338, and 20.1% of Houstonians lived at or below the poverty line.\n",
            "========================\n",
            "\n",
            "Got output: Houston, as of the 2020 U.S. census, has a population of 2,304,580. The city is known for its diverse demographic makeup, with a significant number of undocumented immigrants, who made up nearly 9% of the city's metropolitan population in 2017. The age distribution is quite varied, with a large number of residents under 15 and between the ages of 20 to 34. The median age is 33.4. In terms of housing, 42.3% of Houstonians are homeowners. The median household income in 2019 was $52,338, and 20.1% of the population lived at or below the poverty line.\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Chicago with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Added user message to memory: demographics\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"demographics\"\n",
            "}\n",
            "Got output: Chicago experienced rapid population growth during its first hundred years, becoming one of the fastest-growing cities in the world. From its founding in 1833 with fewer than 200 people, the population grew to over 4,000 within seven years. By 1890, the population had surpassed 1 million, making Chicago the fifth-largest city in the world at the time. The city's population continued to grow, reaching its highest recorded population of 3.6 million in 1950. However, in the latter half of the 20th century, Chicago's population declined, dropping to under 2.7 million by 2010. The city experienced a rise in population for the 2000 census, followed by a decrease in 2010, and then another increase for the 2020 census. According to U.S. census estimates as of July 2019, the largest racial or ethnic groups in Chicago are non-Hispanic White (32.8%), Blacks (30.1%), and Hispanics (29.0%). Additionally, Chicago has the third-largest LGBT population in the United States, with an estimated 7.5% of the adult population identifying as LGBTQ in 2018.\n",
            "========================\n",
            "\n",
            "Got output: Chicago has seen significant population growth since its founding in 1833, with the population reaching over 4,000 within seven years. By 1890, the population had surpassed 1 million, making Chicago the fifth-largest city in the world at the time. The city's population peaked at 3.6 million in 1950, but has since declined to under 2.7 million by 2010. However, there has been a slight increase in population according to the 2020 census.\n",
            "\n",
            "As of July 2019, the largest racial or ethnic groups in Chicago are non-Hispanic White (32.8%), Blacks (30.1%), and Hispanics (29.0%). Additionally, Chicago has the third-largest LGBT population in the United States, with an estimated 7.5% of the adult population identifying as LGBTQ in 2018.\n",
            "========================\n",
            "\n",
            "STARTING TURN 3\n",
            "---------------\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = top_agent.query(\n",
        "    \"Tell the demographics of Houston, and then compare that with the\"\n",
        "    \" demographics of Chicago\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cfe1dd4c-8bfd-43d0-99bc-ca60861dc418",
      "metadata": {
        "id": "cfe1dd4c-8bfd-43d0-99bc-ca60861dc418",
        "outputId": "c9c95af6-2b18-41cb-fd8b-4444fb1c51d4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Houston has a population of 2,304,580, while Chicago has a population of under 2.7 million. Houston is known for its diverse demographic makeup, with a significant number of undocumented immigrants. In terms of age distribution, Houston has a large number of residents under 15 and between the ages of 20 to 34, with a median age of 33.4. In Chicago, the largest racial or ethnic groups are non-Hispanic White, Blacks, and Hispanics. Additionally, Chicago has a significant LGBTQ population, with an estimated 7.5% of the adult population identifying as LGBTQ.\n"
          ]
        }
      ],
      "source": [
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "843c79cb-fc7a-4eac-ba7f-6d25c7a8f08e",
      "metadata": {
        "id": "843c79cb-fc7a-4eac-ba7f-6d25c7a8f08e",
        "outputId": "2def895e-81a0-4b07-c458-446757a8e980",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Houston is a majority-minority city with a diverse population. According to the 2019 U.S. Census Bureau data, the demographics of Houston are as follows: non-Hispanic whites make up 23.3% of the population, Hispanics and Latino Americans make up 45.8%, Blacks or African Americans make up 22.4%, and Asian Americans make up 6.5%. The largest Hispanic or Latino American ethnic group in Houston is Mexican Americans, comprising 31.6% of the population.\n",
            "\n",
            "In comparison, the demographics of Chicago, based on the 2019 U.S. Census Bureau data, are as follows: non-Hispanic whites make up 32.7% of the population, Hispanics and Latino Americans make up 29.9%, Blacks or African Americans make up 29.0%, and Asian Americans make up 7.6%. The largest Hispanic or Latino American ethnic group in Chicago is Mexican Americans, comprising 21.0% of the population.\n",
            "\n",
            "Overall, both Houston and Chicago have diverse populations with significant Hispanic and Latino American, Black or African American, and Asian American communities. However, Houston has a higher percentage of Hispanics and Latino Americans compared to Chicago, while Chicago has a higher percentage of non-Hispanic whites and Blacks or African Americans compared to Houston.\n"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell the demographics of Houston, and then compare that with the\"\n",
        "    \" demographics of Chicago\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e6f99183-bdd9-4d49-8980-44919cf86c03",
      "metadata": {
        "id": "e6f99183-bdd9-4d49-8980-44919cf86c03",
        "outputId": "79cfe915-cf2f-4388-b283-0c4619c544e1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\"=== Early 21st century ===\\nHouston has continued to grow into the 21st century, with the population increasing 15.7% from 2000 to 2022.Oil & gas have continued to fuel Houston's economic growth, with major oil companies including Phillips 66, ConocoPhillips, Occidental Petroleum, Halliburton, and ExxonMobil having their headquarters in the Houston area. In 2001, Enron Corporation, a Houston company with $100 billion in revenue, became engulfed in an accounting scandal which bankrupted the company in 2001. Health care has emerged as a major industry in Houston. The Texas Medical Center is now the largest medical complex in the world and employs over 120,000 people.Three new sports stadiums opened downtown in the first decade of the 21st century. In 2000, the Houston Astros opened their new baseball stadium, Minute Maid Park, in downtown adjacent to the old Union Station. The Houston Texans were formed in 2002 as an NFL expansion team, replacing the Houston Oilers, which had left the city in 1996. NRG Stadium opened the same year. In 2003, the Toyota Center opened as the home for the Houston Rockets. In 2005, the Houston Dynamo soccer team was formed. In 2017, the Houston Astros won their first World Series.\\n\\nFlooding has been a recurring problem in the Houston area, exacerbated by a lack of zoning laws, which allowed unregulated building of residential homes and other structures in flood-prone areas. In June 2001, Tropical Storm Allison dumped up to 40 inches (1,000 mm) of rain on parts of Houston, causing what was then the worst flooding in the city's history and billions of dollars in damage, and killed 20 people in Texas. In August 2005, Houston became a shelter to more than 150,000 people from New Orleans, who evacuated from Hurricane Katrina. One month later, about 2.5 million Houston-area residents evacuated when Hurricane Rita approached the Gulf Coast, leaving little damage to the Houston area. This was the largest urban evacuation in the history of the United States. In May 2015, seven people died after 12 inches of rain fell in 10 hours during what is known as the Memorial Day Flood. Eight people died in April 2016 during a storm that dropped 17 inches of rain. The worst came in late August 2017, when Hurricane Harvey stalled over southeastern Texas, much like Tropical Storm Allison did sixteen years earlier, causing severe flooding in the Houston area, with some areas receiving over 50 inches (1,300 mm) of rain. The rainfall exceeded 50 inches in several areas locally, breaking the national record for rainfall. The damage for the Houston area was estimated at up to $125 billion U.S. dollars, and was considered to be one of the worst natural disasters in the history of the United States, with the death toll exceeding 70 people.\\n\\n\\n== Geography ==\\n\\nHouston is 165 miles (266 km) east of Austin, 88 miles (142 km) west of the Louisiana border, and 250 miles (400 km) south of Dallas. The city has a total area of 637.4 square miles (1,651 km2); this comprises over 599.59 square miles (1,552.9 km2) of land and 22.3 square miles (58 km2) covered by water. Most of Houston is on the gulf coastal plain, and its vegetation is classified as Western Gulf coastal grasslands while further north, it transitions into a subtropical jungle, the Big Thicket.\\nMuch of the city was built on forested land, marshes, or swamps, and all are still visible in surrounding areas. Flat terrain and extensive greenfield development have combined to worsen flooding. Downtown stands about 50 feet (15 m) above sea level, and the highest point in far northwest Houston is about 150 feet (46 m) in elevation. The city once relied on groundwater for its needs, but land subsidence forced the city to turn to ground-level water sources such as Lake Houston, Lake Conroe, and Lake Livingston. The city owns surface water rights for 1.20 billion US gallons (4.5 Gl) of water a day in addition to 150 million US gallons (570 Ml) a day of groundwater.Houston has four major bayous passing through the city that accept water from the extensive drainage system. Buffalo Bayou runs through Downtown and the Houston Ship Channel, and has three tributaries: White Oak Bayou, which runs through the Houston Heights community northwest of Downtown and then towards Downtown; Brays Bayou, which runs along the Texas Medical Center; and Sims Bayou, which runs through the south of Houston and Downtown Houston. The ship channel continues past Galveston and then into the Gulf of Mexico.\""
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 25
        }
      ],
      "source": [
        "# baseline: the response tells you nothing about Chicago...\n",
        "response.source_nodes[3].get_content()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5b7e0a4c-b7be-4797-9c82-1577693bd117",
      "metadata": {
        "id": "5b7e0a4c-b7be-4797-9c82-1577693bd117",
        "outputId": "57a2a7af-a23c-48a5-ad50-be0170a60c1b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Shanghai with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Added user message to memory: history\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Got output: Shanghai has a rich history that dates back to the 19th century. During this time, the city gained international attention due to its economic and trade potential. It became one of the five treaty ports for international trade after the First Opium War. Shanghai experienced various conflicts and attacks, including the Taiping Rebellion and the First Sino-Japanese War. In the early 20th century, Shanghai became a major commercial and financial hub, attracting people from all over the world. However, the city also faced challenges during the Cultural Revolution. Since then, Shanghai has undergone significant economic reforms and development, becoming a global city and a major contributor to China's tax revenue.\n",
            "========================\n",
            "\n",
            "Got output: Shanghai has a rich history that dates back to the 19th century. It gained international attention due to its economic and trade potential and became one of the five treaty ports for international trade after the First Opium War. The city experienced various conflicts and attacks, including the Taiping Rebellion and the First Sino-Japanese War. \n",
            "\n",
            "In the early 20th century, Shanghai became a major commercial and financial hub, attracting people from all over the world. However, the city also faced challenges during the Cultural Revolution. Since then, Shanghai has undergone significant economic reforms and development, becoming a global city and a major contributor to China's tax revenue.\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Beijing with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Added user message to memory: history\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"history\"\n",
            "}\n",
            "Got output: Beijing has a rich history that spans several dynasties and significant events. It was the capital of the People's Republic of China, which was established in 1949. During the Cultural Revolution from 1966 to 1976, Beijing was the center of the Red Guard movement and experienced political purges. In recent years, the city has faced challenges related to urbanization, such as heavy traffic and poor air quality. Beijing has also hosted major international events, including the 2008 Summer Olympics and the upcoming 2022 Winter Olympics. Prior to the People's Republic of China, Beijing was the capital during the Qing dynasty and the Republic of China. It played a role in the Second Opium War and the Boxer Rebellion. The city has seen changes in its name, from Beijing to Beiping and back to Beijing.\n",
            "========================\n",
            "\n",
            "Got output: Beijing has a rich and complex history that spans several dynasties and significant events. It was the capital of the People's Republic of China, which was established in 1949. During the Cultural Revolution from 1966 to 1976, Beijing was the center of the Red Guard movement and experienced political purges. \n",
            "\n",
            "In recent years, the city has faced challenges related to urbanization, such as heavy traffic and poor air quality. Despite these challenges, Beijing has also hosted major international events, including the 2008 Summer Olympics and the upcoming 2022 Winter Olympics.\n",
            "\n",
            "Before the establishment of the People's Republic of China, Beijing was the capital during the Qing dynasty and the Republic of China. It played a significant role in the Second Opium War and the Boxer Rebellion. The city has also seen changes in its name, from Beijing to Beiping and back to Beijing.\n",
            "========================\n",
            "\n",
            "STARTING TURN 3\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Shanghai with args: {\n",
            "  \"input\": \"economy\"\n",
            "}\n",
            "Added user message to memory: economy\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"economy\"\n",
            "}\n",
            "Got output: Shanghai has a booming economy and is considered a global center for finance and innovation. It is a national center for commerce, trade, and transportation, with the world's busiest container port. The Greater Shanghai metropolitan area, which includes several other cities, has a significant gross metropolitan product. Shanghai's economy is driven by industries such as retail, finance, IT, real estate, machine manufacturing, and automotive manufacturing. The city has experienced rapid development and has been one of the fastest-growing cities in the world. It is also known for hosting the China International Import Expo. Shanghai's economy is projected to continue growing, and it is ranked among the top ten largest metropolitan economies in the world.\n",
            "========================\n",
            "\n",
            "Got output: Shanghai's economy is thriving and globally recognized for its finance and innovation sectors. It serves as a national hub for commerce, trade, and transportation, boasting the world's busiest container port. The Greater Shanghai metropolitan area, inclusive of several other cities, contributes significantly to the gross metropolitan product. Key industries driving Shanghai's economy include retail, finance, IT, real estate, machine manufacturing, and automotive manufacturing. The city has witnessed rapid development and is one of the fastest-growing cities globally. It also hosts the China International Import Expo. With a projected continuous growth, Shanghai ranks among the top ten largest metropolitan economies worldwide.\n",
            "========================\n",
            "\n",
            "STARTING TURN 4\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: tool_Beijing with args: {\n",
            "  \"input\": \"economy\"\n",
            "}\n",
            "Added user message to memory: economy\n",
            "=== Calling Function ===\n",
            "Calling function: vector_tool with args: {\n",
            "  \"input\": \"economy\"\n",
            "}\n",
            "Got output: Beijing has a post-industrial economy that is dominated by the tertiary sector (services), which generates the majority of the city's output. The services sector is diversified, with professional services, wholesale and retail, information technology, commercial real estate, scientific research, and residential real estate contributing significantly to the city's economy. The secondary sector (manufacturing, construction) and the primary sector (agriculture, mining) also contribute to the economy, although to a lesser extent. Beijing has several high-end economic output zones that drive local economic growth, including Zhongguancun, Beijing Financial Street, Beijing Central Business District (CBD), Beijing Economic and Technological Development Area (Yizhuang), Beijing Airport Economic Zone, and Beijing Olympic Center Zone. These zones are home to various industries and sectors, such as technology companies, financial institutions, professional service firms, industrial parks, and automobile assembly plants. Additionally, Beijing has been described as the \"billionaire capital of the world\" and ranks among the top major cities globally in terms of influence and competitiveness.\n",
            "========================\n",
            "\n",
            "Got output: Beijing's economy is primarily driven by the tertiary sector, which includes services such as professional services, wholesale and retail, information technology, commercial real estate, scientific research, and residential real estate. The secondary sector, which includes manufacturing and construction, and the primary sector, which includes agriculture and mining, also contribute to the economy but to a lesser extent.\n",
            "\n",
            "The city has several high-end economic output zones that drive local economic growth. These include Zhongguancun, Beijing Financial Street, Beijing Central Business District (CBD), Beijing Economic and Technological Development Area (Yizhuang), Beijing Airport Economic Zone, and Beijing Olympic Center Zone. These zones host a variety of industries and sectors, such as technology companies, financial institutions, professional service firms, industrial parks, and automobile assembly plants.\n",
            "\n",
            "Beijing is often referred to as the \"billionaire capital of the world\" and is considered one of the top major cities globally in terms of influence and competitiveness.\n",
            "========================\n",
            "\n",
            "STARTING TURN 5\n",
            "---------------\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = top_agent.query(\n",
        "    \"Tell me the differences between Shanghai and Beijing in terms of history\"\n",
        "    \" and current economy\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "98540ec5-093d-41a0-888b-246ee7093cff",
      "metadata": {
        "id": "98540ec5-093d-41a0-888b-246ee7093cff",
        "outputId": "b4b14539-5591-4988-e71c-90f9ab092827",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "In terms of history, both Shanghai and Beijing have rich and significant pasts. Shanghai gained international attention in the 19th century as a major trade port and has experienced conflicts and developments throughout the years. Beijing, on the other hand, has been the capital of China for several dynasties and has witnessed significant events such as the Cultural Revolution. \n",
            "\n",
            "In terms of current economy, Shanghai is known for its thriving economy, particularly in finance and innovation sectors. It serves as a national hub for commerce, trade, and transportation and is one of the fastest-growing cities globally. Key industries in Shanghai include retail, finance, IT, real estate, and manufacturing. \n",
            "\n",
            "Beijing's economy is primarily driven by the tertiary sector, with a focus on services such as professional services, information technology, and commercial real estate. It also has high-end economic output zones that contribute to its growth. Beijing is considered one of the top major cities globally in terms of influence and competitiveness.\n",
            "\n",
            "Overall, while both cities have significant historical backgrounds, Shanghai is known for its economic prowess and global recognition, while Beijing serves as the political and cultural capital of China.\n"
          ]
        }
      ],
      "source": [
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a0f49e28-e85e-4c93-8eed-68122ad9f537",
      "metadata": {
        "id": "a0f49e28-e85e-4c93-8eed-68122ad9f537",
        "outputId": "cd37a5b9-0d97-44fe-cdb9-c78187f68407",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 245
        }
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'base_query_engine' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-2-dec7fc0fb015>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# baseline\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m response = base_query_engine.query(\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;34m\"Tell me the differences between Shanghai and Beijing in terms of history\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0;34m\" and current economy\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m )\n",
            "\u001b[0;31mNameError\u001b[0m: name 'base_query_engine' is not defined"
          ]
        }
      ],
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"Tell me the differences between Shanghai and Beijing in terms of history\"\n",
        "    \" and current economy\"\n",
        ")\n",
        "print(str(response))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# baseline\n",
        "response = base_query_engine.query(\n",
        "    \"List ten differences between Toronto and Boston\"\n",
        ")\n",
        "print(str(response))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 245
        },
        "id": "sXSaBMOKBiEz",
        "outputId": "fbe6fd97-2221-46c6-850e-edf349cba2c4"
      },
      "id": "sXSaBMOKBiEz",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'base_query_engine' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-d41eff9f2e9c>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# baseline\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m response = base_query_engine.query(\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;34m\"List ten differences between Toronto and Boston\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m )\n\u001b[1;32m      5\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresponse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'base_query_engine' is not defined"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "llama_index_v2",
      "language": "python",
      "name": "llama_index_v2"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}