Module_03_LlamaIndex_RAG_AGENT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ashispapu/LLMs/blob/main/Module_03_LlamaIndex_RAG_AGENT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q llama-index==0.9.14.post3 deeplake==3.8.8 openai==1.3.8 cohere==4.37"
      ],
      "metadata": {
        "id": "oLsz4honE_jq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = '<YOUR_OPENAI_API_KEY>'\n",
        "os.environ['ACTIVELOOP_TOKEN'] = '<YOUR_ACTIVELOOP_API_KEY>'"
      ],
      "metadata": {
        "id": "Uh_M0Z0FFJPb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare Indexes"
      ],
      "metadata": {
        "id": "e4xCp6-mN3xw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir -p 'data/1k/'\n",
        "!wget 'https://github.com/idontcalculate/data-repo/blob/main/machine_to_end_war.txt' -O './data/1k/tesla.txt'\n",
        "!wget 'https://github.com/idontcalculate/data-repo/blob/main/prodigal_chapter10.txt' -O './data/1k/web.txt'"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mE86KzdvOPgX",
        "outputId": "becf09cb-35a4-4c6c-f26a-6cdb41fbacd5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-12-15 16:23:50--  https://github.com/idontcalculate/data-repo/blob/main/machine_to_end_war.txt\n",
            "Resolving github.com (github.com)... 140.82.113.3\n",
            "Connecting to github.com (github.com)|140.82.113.3|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 18616 (18K) [text/plain]\n",
            "Saving to: ‘./data/1k/tesla.txt’\n",
            "\n",
            "./data/1k/tesla.txt 100%[===================>]  18.18K  --.-KB/s    in 0.04s   \n",
            "\n",
            "2023-12-15 16:23:50 (424 KB/s) - ‘./data/1k/tesla.txt’ saved [18616/18616]\n",
            "\n",
            "--2023-12-15 16:23:50--  https://github.com/idontcalculate/data-repo/blob/main/prodigal_chapter10.txt\n",
            "Resolving github.com (github.com)... 140.82.113.4\n",
            "Connecting to github.com (github.com)|140.82.113.4|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 28197 (28K) [text/plain]\n",
            "Saving to: ‘./data/1k/web.txt’\n",
            "\n",
            "./data/1k/web.txt   100%[===================>]  27.54K  --.-KB/s    in 0.04s   \n",
            "\n",
            "2023-12-15 16:23:51 (628 KB/s) - ‘./data/1k/web.txt’ saved [28197/28197]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### From VectorStore"
      ],
      "metadata": {
        "id": "xAXcBI5COHWy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import SimpleDirectoryReader\n",
        "\n",
        "tesla_docs = SimpleDirectoryReader( input_files=[\"/content/data/1k/tesla.txt\"] ).load_data()"
      ],
      "metadata": {
        "id": "ufXfJwfcORAa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.vector_stores import DeepLakeVectorStore\n",
        "\n",
        "my_activeloop_org_id = \"genai360\"\n",
        "my_activeloop_dataset_name = \"LlamaIndex_tesla_predictions\"\n",
        "dataset_path = f\"hub://{my_activeloop_org_id}/{my_activeloop_dataset_name}\"\n",
        "\n",
        "# Create an index over the documnts\n",
        "vector_store = DeepLakeVectorStore(dataset_path=dataset_path, overwrite=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "--NnpT4UOQbv",
        "outputId": "00b40ba0-ae85-412d-c897-8d5ce317b172"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/deeplake/util/check_latest_version.py:32: UserWarning: A newer version of deeplake (3.8.12) is available. It's recommended that you update to the latest version using `pip install -U deeplake`.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your Deep Lake dataset has been successfully created!\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.storage.storage_context import StorageContext\n",
        "\n",
        "storage_context = StorageContext.from_defaults(vector_store=vector_store)"
      ],
      "metadata": {
        "id": "loC5XGShPtf_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import VectorStoreIndex\n",
        "\n",
        "tesla_index = VectorStoreIndex.from_documents(tesla_docs, storage_context=storage_context)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7pPld5m3Ptco",
        "outputId": "8b9f21fd-0780-4d61-8a6e-97d524d68232"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Uploading data to deeplake dataset.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5/5 [00:00<00:00,  7.17it/s]\n",
            "/"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset(path='hub://genai360/LlamaIndex_tesla_predictions', tensors=['text', 'metadata', 'embedding', 'id'])\n",
            "\n",
            "  tensor      htype      shape     dtype  compression\n",
            "  -------    -------    -------   -------  ------- \n",
            "   text       text      (5, 1)      str     None   \n",
            " metadata     json      (5, 1)      str     None   \n",
            " embedding  embedding  (5, 1536)  float32   None   \n",
            "    id        text      (5, 1)      str     None   \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r \r"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## From Local Index"
      ],
      "metadata": {
        "id": "GIQLCDvuQPPU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "webtext_docs = SimpleDirectoryReader( input_files=[\"/content/data/1k/web.txt\"] ).load_data()"
      ],
      "metadata": {
        "id": "7lPWnGE_PtaR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  # Try to load the index if it is already calculated\n",
        "  storage_context = StorageContext.from_defaults( persist_dir=\"/content/storage/webtext\" )\n",
        "  webtext_index = load_index_from_storage(storage_context)\n",
        "  print(\"Loaded the pre-computed index.\")\n",
        "except:\n",
        "  # Otherwise, generate the indexes\n",
        "  webtext_index = VectorStoreIndex.from_documents(webtext_docs)\n",
        "  webtext_index.storage_context.persist(persist_dir=\"/content/storage/webtext\")\n",
        "  print(\"Generated the index.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P1ieq9i8QUa1",
        "outputId": "f9da5d43-fbea-494c-9751-21134dc4156c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated the index.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create Query Enginges"
      ],
      "metadata": {
        "id": "dF3KH4bHRU9L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tesla_engine = tesla_index.as_query_engine(similarity_top_k=3)\n",
        "webtext_engine = webtext_index.as_query_engine(similarity_top_k=3)"
      ],
      "metadata": {
        "id": "MR4BXAVqQnV8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create the Tools"
      ],
      "metadata": {
        "id": "YFbYTky1Rgu0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.tools import QueryEngineTool, ToolMetadata\n",
        "\n",
        "query_engine_tools = [\n",
        "    QueryEngineTool(\n",
        "        query_engine=tesla_engine,\n",
        "        metadata=ToolMetadata(\n",
        "            name=\"tesla_1k\",\n",
        "            description=(\n",
        "                \"Provides information about Tesla's statements that refers to future times and predictions. \"\n",
        "                \"Use a detailed plain text question as input to the tool.\"\n",
        "            ),\n",
        "        ),\n",
        "    ),\n",
        "    QueryEngineTool(\n",
        "        query_engine=webtext_engine,\n",
        "        metadata=ToolMetadata(\n",
        "            name=\"webtext_1k\",\n",
        "            description=(\n",
        "                \"Provides information about tesla's life and biographical data. \"\n",
        "                \"Use a detailed plain text question as input to the tool.\"\n",
        "            ),\n",
        "        ),\n",
        "    ),\n",
        "]"
      ],
      "metadata": {
        "id": "vMptZI0uPtVU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Define the Agent"
      ],
      "metadata": {
        "id": "0CU3MaJ4Rt4R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.agent import OpenAIAgent\n",
        "agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)"
      ],
      "metadata": {
        "id": "RdBtZi50PtSe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "agent.chat_repl()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M-0kHgKCRiq0",
        "outputId": "8ab2c744-e818-46ec-be40-da17fb1d11f0"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "===== Entering Chat REPL =====\n",
            "Type \"exit\" to exit.\n",
            "\n",
            "Human: What influenced Nikola Tesla to become an inventor?\n",
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: webtext_1k with args: {\n",
            "\"input\": \"What influenced Nikola Tesla to become an inventor?\"\n",
            "}\n",
            "Got output: Nikola Tesla was influenced to become an inventor by his studies of mechanical vibrations. He observed the selective response of objects to vibrations and realized the potential for producing effects of tremendous magnitude on physical objects. This led him to pursue research in the field of high-frequency and high-potential currents, which eventually resulted in his groundbreaking inventions.\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n",
            "Assistant: Nikola Tesla was influenced to become an inventor by his studies of mechanical vibrations. He observed the selective response of objects to vibrations and realized the potential for producing effects of tremendous magnitude on physical objects. This led him to pursue research in the field of high-frequency and high-potential currents, which eventually resulted in his groundbreaking inventions.\n",
            "\n",
            "Human: exit\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Agents with Tools"
      ],
      "metadata": {
        "id": "AGakE3sO1F_5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.tools import FunctionTool\n",
        "\n",
        "def multiply(a: int, b: int) -> int:\n",
        "    \"\"\"Multiply 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",
        "multiply_tool = FunctionTool.from_defaults(fn=multiply, name=\"multiply\")\n",
        "add_tool = FunctionTool.from_defaults(fn=add, name=\"add\")\n",
        "\n",
        "all_tools = [multiply_tool, add_tool]"
      ],
      "metadata": {
        "id": "xwTwYSSjRihM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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",
        ")"
      ],
      "metadata": {
        "id": "gRNaPGDyRieS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.agent import FnRetrieverOpenAIAgent\n",
        "\n",
        "agent = FnRetrieverOpenAIAgent.from_retriever(\n",
        "    obj_index.as_retriever(), verbose=True\n",
        ")"
      ],
      "metadata": {
        "id": "GNHxglUDKJ5L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "agent.chat(\"What's 12 multiplied by 22? Make sure to use Tools\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uSCd47OB2ZIN",
        "outputId": "665317f3-6710-4a72-f940-96fd50800ef0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: multiply with args: {\n",
            "  \"a\": 12,\n",
            "  \"b\": 22\n",
            "}\n",
            "Got output: 264\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "AgentChatResponse(response='12 multiplied by 22 is 264.', sources=[ToolOutput(content='264', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 12, 'b': 22}}, raw_output=264)], source_nodes=[])"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "agent.chat( \"What is 5 + 2?\", tool_choice=\"add\" )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "btL4lvUy2ZBK",
        "outputId": "6b36a503-321e-489f-8101-e20890367ddb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "STARTING TURN 1\n",
            "---------------\n",
            "\n",
            "=== Calling Function ===\n",
            "Calling function: add with args: {\n",
            "  \"a\": 5,\n",
            "  \"b\": 2\n",
            "}\n",
            "Got output: 7\n",
            "========================\n",
            "\n",
            "STARTING TURN 2\n",
            "---------------\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "AgentChatResponse(response='5 + 2 is equal to 7.', sources=[ToolOutput(content='7', tool_name='add', raw_input={'args': (), 'kwargs': {'a': 5, 'b': 2}}, raw_output=7)], source_nodes=[])"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "GHfXq66zEZ7B"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}