python/llamaindex-integration/RAG LlamaIndex Workflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b078478",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: llama-index in /opt/homebrew/lib/python3.12/site-packages (0.10.57)\n",
      "Collecting llama-index\n",
      "  Downloading llama_index-0.11.9-py3-none-any.whl.metadata (11 kB)\n",
      "Collecting llama-index-agent-openai<0.4.0,>=0.3.1 (from llama-index)\n",
      "  Downloading llama_index_agent_openai-0.3.1-py3-none-any.whl.metadata (677 bytes)\n",
      "Collecting llama-index-cli<0.4.0,>=0.3.1 (from llama-index)\n",
      "  Downloading llama_index_cli-0.3.1-py3-none-any.whl.metadata (1.5 kB)\n",
      "Collecting llama-index-core<0.12.0,>=0.11.9 (from llama-index)\n",
      "  Downloading llama_index_core-0.11.9-py3-none-any.whl.metadata (2.4 kB)\n",
      "Collecting llama-index-embeddings-openai<0.3.0,>=0.2.4 (from llama-index)\n",
      "  Downloading llama_index_embeddings_openai-0.2.4-py3-none-any.whl.metadata (635 bytes)\n",
      "Collecting llama-index-indices-managed-llama-cloud>=0.3.0 (from llama-index)\n",
      "  Downloading llama_index_indices_managed_llama_cloud-0.3.0-py3-none-any.whl.metadata (3.8 kB)\n",
      "Requirement already satisfied: llama-index-legacy<0.10.0,>=0.9.48 in /opt/homebrew/lib/python3.12/site-packages (from llama-index) (0.9.48)\n",
      "Collecting llama-index-llms-openai<0.3.0,>=0.2.3 (from llama-index)\n",
      "  Downloading llama_index_llms_openai-0.2.5-py3-none-any.whl.metadata (705 bytes)\n",
      "Collecting llama-index-multi-modal-llms-openai<0.3.0,>=0.2.0 (from llama-index)\n",
      "  Downloading llama_index_multi_modal_llms_openai-0.2.0-py3-none-any.whl.metadata (728 bytes)\n",
      "Collecting llama-index-program-openai<0.3.0,>=0.2.0 (from llama-index)\n",
      "  Downloading llama_index_program_openai-0.2.0-py3-none-any.whl.metadata (766 bytes)\n",
      "Collecting llama-index-question-gen-openai<0.3.0,>=0.2.0 (from llama-index)\n",
      "  Downloading llama_index_question_gen_openai-0.2.0-py3-none-any.whl.metadata (785 bytes)\n",
      "Collecting llama-index-readers-file<0.3.0,>=0.2.0 (from llama-index)\n",
      "  Downloading llama_index_readers_file-0.2.1-py3-none-any.whl.metadata (5.4 kB)\n",
      "Collecting llama-index-readers-llama-parse>=0.3.0 (from llama-index)\n",
      "  Downloading llama_index_readers_llama_parse-0.3.0-py3-none-any.whl.metadata (3.5 kB)\n",
      "Collecting nltk>3.8.1 (from llama-index)\n",
      "  Downloading nltk-3.9.1-py3-none-any.whl.metadata (2.9 kB)\n",
      "Requirement already satisfied: openai>=1.14.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-agent-openai<0.4.0,>=0.3.1->llama-index) (1.44.1)\n",
      "Requirement already satisfied: PyYAML>=6.0.1 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (6.0.1)\n",
      "Requirement already satisfied: SQLAlchemy>=1.4.49 in /opt/homebrew/lib/python3.12/site-packages (from SQLAlchemy[asyncio]>=1.4.49->llama-index-core<0.12.0,>=0.11.9->llama-index) (2.0.28)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.6 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (3.9.3)\n",
      "Requirement already satisfied: dataclasses-json in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (0.5.14)\n",
      "Requirement already satisfied: deprecated>=1.2.9.3 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (1.2.14)\n",
      "Requirement already satisfied: dirtyjson<2.0.0,>=1.0.8 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (1.0.8)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (2024.2.0)\n",
      "Requirement already satisfied: httpx in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (0.27.0)\n",
      "Requirement already satisfied: nest-asyncio<2.0.0,>=1.5.8 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (1.6.0)\n",
      "Requirement already satisfied: networkx>=3.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (3.3)\n",
      "Requirement already satisfied: numpy<2.0.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (1.26.4)\n",
      "Requirement already satisfied: pillow>=9.0.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (10.4.0)\n",
      "Requirement already satisfied: pydantic<3.0.0,>=2.7.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (2.8.2)\n",
      "Requirement already satisfied: requests>=2.31.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (2.31.0)\n",
      "Requirement already satisfied: tenacity!=8.4.0,<9.0.0,>=8.2.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (8.2.3)\n",
      "Requirement already satisfied: tiktoken>=0.3.3 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (0.7.0)\n",
      "Requirement already satisfied: tqdm<5.0.0,>=4.66.1 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (4.66.2)\n",
      "Requirement already satisfied: typing-extensions>=4.5.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (4.12.2)\n",
      "Requirement already satisfied: typing-inspect>=0.8.0 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (0.9.0)\n",
      "Requirement already satisfied: wrapt in /opt/homebrew/lib/python3.12/site-packages (from llama-index-core<0.12.0,>=0.11.9->llama-index) (1.16.0)\n",
      "Collecting llama-cloud>=0.0.11 (from llama-index-indices-managed-llama-cloud>=0.3.0->llama-index)\n",
      "  Downloading llama_cloud-0.0.17-py3-none-any.whl.metadata (751 bytes)\n",
      "Requirement already satisfied: pandas in /opt/homebrew/lib/python3.12/site-packages (from llama-index-legacy<0.10.0,>=0.9.48->llama-index) (2.2.1)\n",
      "Requirement already satisfied: beautifulsoup4<5.0.0,>=4.12.3 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-readers-file<0.3.0,>=0.2.0->llama-index) (4.12.3)\n",
      "Requirement already satisfied: pypdf<5.0.0,>=4.0.1 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-readers-file<0.3.0,>=0.2.0->llama-index) (4.3.1)\n",
      "Requirement already satisfied: striprtf<0.0.27,>=0.0.26 in /opt/homebrew/lib/python3.12/site-packages (from llama-index-readers-file<0.3.0,>=0.2.0->llama-index) (0.0.26)\n",
      "Collecting llama-parse>=0.5.0 (from llama-index-readers-llama-parse>=0.3.0->llama-index)\n",
      "  Downloading llama_parse-0.5.5-py3-none-any.whl.metadata (6.1 kB)\n",
      "Requirement already satisfied: click in /opt/homebrew/lib/python3.12/site-packages (from nltk>3.8.1->llama-index) (8.1.7)\n",
      "Requirement already satisfied: joblib in /opt/homebrew/lib/python3.12/site-packages (from nltk>3.8.1->llama-index) (1.4.2)\n",
      "Requirement already satisfied: regex>=2021.8.3 in /opt/homebrew/lib/python3.12/site-packages (from nltk>3.8.1->llama-index) (2023.12.25)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/homebrew/lib/python3.12/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/homebrew/lib/python3.12/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.12.0,>=0.11.9->llama-index) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/homebrew/lib/python3.12/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/homebrew/lib/python3.12/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.12.0,>=0.11.9->llama-index) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/homebrew/lib/python3.12/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.9.4)\n",
      "Requirement already satisfied: soupsieve>1.2 in /opt/homebrew/lib/python3.12/site-packages (from beautifulsoup4<5.0.0,>=4.12.3->llama-index-readers-file<0.3.0,>=0.2.0->llama-index) (2.5)\n",
      "Requirement already satisfied: anyio in /opt/homebrew/lib/python3.12/site-packages (from httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (3.7.1)\n",
      "Requirement already satisfied: certifi in /opt/homebrew/lib/python3.12/site-packages (from httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (2024.2.2)\n",
      "Requirement already satisfied: httpcore==1.* in /opt/homebrew/lib/python3.12/site-packages (from httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.0.5)\n",
      "Requirement already satisfied: idna in /opt/homebrew/lib/python3.12/site-packages (from httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (3.6)\n",
      "Requirement already satisfied: sniffio in /opt/homebrew/lib/python3.12/site-packages (from httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.3.1)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /opt/homebrew/lib/python3.12/site-packages (from httpcore==1.*->httpx->llama-index-core<0.12.0,>=0.11.9->llama-index) (0.14.0)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /opt/homebrew/lib/python3.12/site-packages (from openai>=1.14.0->llama-index-agent-openai<0.4.0,>=0.3.1->llama-index) (1.9.0)\n",
      "Requirement already satisfied: jiter<1,>=0.4.0 in /opt/homebrew/lib/python3.12/site-packages (from openai>=1.14.0->llama-index-agent-openai<0.4.0,>=0.3.1->llama-index) (0.5.0)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /opt/homebrew/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.0->llama-index-core<0.12.0,>=0.11.9->llama-index) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.20.1 in /opt/homebrew/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.0->llama-index-core<0.12.0,>=0.11.9->llama-index) (2.20.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/homebrew/lib/python3.12/site-packages (from requests>=2.31.0->llama-index-core<0.12.0,>=0.11.9->llama-index) (3.3.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/homebrew/lib/python3.12/site-packages (from requests>=2.31.0->llama-index-core<0.12.0,>=0.11.9->llama-index) (2.2.1)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /opt/homebrew/lib/python3.12/site-packages (from SQLAlchemy[asyncio]>=1.4.49->llama-index-core<0.12.0,>=0.11.9->llama-index) (3.0.3)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /opt/homebrew/lib/python3.12/site-packages (from typing-inspect>=0.8.0->llama-index-core<0.12.0,>=0.11.9->llama-index) (1.0.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /opt/homebrew/lib/python3.12/site-packages (from dataclasses-json->llama-index-core<0.12.0,>=0.11.9->llama-index) (3.21.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /opt/homebrew/lib/python3.12/site-packages (from pandas->llama-index-legacy<0.10.0,>=0.9.48->llama-index) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/homebrew/lib/python3.12/site-packages (from pandas->llama-index-legacy<0.10.0,>=0.9.48->llama-index) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /opt/homebrew/lib/python3.12/site-packages (from pandas->llama-index-legacy<0.10.0,>=0.9.48->llama-index) (2024.1)\n",
      "Requirement already satisfied: packaging>=17.0 in /opt/homebrew/lib/python3.12/site-packages (from marshmallow<4.0.0,>=3.18.0->dataclasses-json->llama-index-core<0.12.0,>=0.11.9->llama-index) (23.2)\n",
      "Requirement already satisfied: six>=1.5 in /opt/homebrew/lib/python3.12/site-packages (from python-dateutil>=2.8.2->pandas->llama-index-legacy<0.10.0,>=0.9.48->llama-index) (1.16.0)\n",
      "Downloading llama_index-0.11.9-py3-none-any.whl (6.8 kB)\n",
      "Downloading llama_index_agent_openai-0.3.1-py3-none-any.whl (13 kB)\n",
      "Downloading llama_index_cli-0.3.1-py3-none-any.whl (27 kB)\n",
      "Downloading llama_index_core-0.11.9-py3-none-any.whl (1.6 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m33.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading llama_index_embeddings_openai-0.2.4-py3-none-any.whl (6.1 kB)\n",
      "Downloading llama_index_indices_managed_llama_cloud-0.3.0-py3-none-any.whl (9.5 kB)\n",
      "Downloading llama_index_llms_openai-0.2.5-py3-none-any.whl (12 kB)\n",
      "Downloading llama_index_multi_modal_llms_openai-0.2.0-py3-none-any.whl (5.9 kB)\n",
      "Downloading llama_index_program_openai-0.2.0-py3-none-any.whl (5.3 kB)\n",
      "Downloading llama_index_question_gen_openai-0.2.0-py3-none-any.whl (2.9 kB)\n",
      "Downloading llama_index_readers_file-0.2.1-py3-none-any.whl (38 kB)\n",
      "Downloading llama_index_readers_llama_parse-0.3.0-py3-none-any.whl (2.5 kB)\n",
      "Downloading nltk-3.9.1-py3-none-any.whl (1.5 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m33.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading llama_cloud-0.0.17-py3-none-any.whl (187 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m187.4/187.4 kB\u001b[0m \u001b[31m21.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading llama_parse-0.5.5-py3-none-any.whl (10 kB)\n",
      "Installing collected packages: nltk, llama-index-core, llama-cloud, llama-parse, llama-index-readers-file, llama-index-indices-managed-llama-cloud, llama-index-embeddings-openai, llama-index-readers-llama-parse, llama-index-llms-openai, llama-index-agent-openai, llama-index-program-openai, llama-index-question-gen-openai, llama-index-multi-modal-llms-openai, llama-index-cli, llama-index\n",
      "  Attempting uninstall: nltk\n",
      "    Found existing installation: nltk 3.8.1\n",
      "    Uninstalling nltk-3.8.1:\n",
      "      Successfully uninstalled nltk-3.8.1\n",
      "  Attempting uninstall: llama-index-core\n",
      "    Found existing installation: llama-index-core 0.10.57\n",
      "    Uninstalling llama-index-core-0.10.57:\n",
      "      Successfully uninstalled llama-index-core-0.10.57\n",
      "  Attempting uninstall: llama-cloud\n",
      "    Found existing installation: llama-cloud 0.0.10\n",
      "    Uninstalling llama-cloud-0.0.10:\n",
      "      Successfully uninstalled llama-cloud-0.0.10\n",
      "  Attempting uninstall: llama-parse\n",
      "    Found existing installation: llama-parse 0.4.9\n",
      "    Uninstalling llama-parse-0.4.9:\n",
      "      Successfully uninstalled llama-parse-0.4.9\n",
      "  Attempting uninstall: llama-index-readers-file\n",
      "    Found existing installation: llama-index-readers-file 0.1.30\n",
      "    Uninstalling llama-index-readers-file-0.1.30:\n",
      "      Successfully uninstalled llama-index-readers-file-0.1.30\n",
      "  Attempting uninstall: llama-index-indices-managed-llama-cloud\n",
      "    Found existing installation: llama-index-indices-managed-llama-cloud 0.2.5\n",
      "    Uninstalling llama-index-indices-managed-llama-cloud-0.2.5:\n",
      "      Successfully uninstalled llama-index-indices-managed-llama-cloud-0.2.5\n",
      "  Attempting uninstall: llama-index-embeddings-openai\n",
      "    Found existing installation: llama-index-embeddings-openai 0.1.11\n",
      "    Uninstalling llama-index-embeddings-openai-0.1.11:\n",
      "      Successfully uninstalled llama-index-embeddings-openai-0.1.11\n",
      "  Attempting uninstall: llama-index-readers-llama-parse\n",
      "    Found existing installation: llama-index-readers-llama-parse 0.1.6\n",
      "    Uninstalling llama-index-readers-llama-parse-0.1.6:\n",
      "      Successfully uninstalled llama-index-readers-llama-parse-0.1.6\n",
      "  Attempting uninstall: llama-index-llms-openai\n",
      "    Found existing installation: llama-index-llms-openai 0.1.27\n",
      "    Uninstalling llama-index-llms-openai-0.1.27:\n",
      "      Successfully uninstalled llama-index-llms-openai-0.1.27\n",
      "  Attempting uninstall: llama-index-agent-openai\n",
      "    Found existing installation: llama-index-agent-openai 0.2.9\n",
      "    Uninstalling llama-index-agent-openai-0.2.9:\n",
      "      Successfully uninstalled llama-index-agent-openai-0.2.9\n",
      "  Attempting uninstall: llama-index-program-openai\n",
      "    Found existing installation: llama-index-program-openai 0.1.7\n",
      "    Uninstalling llama-index-program-openai-0.1.7:\n",
      "      Successfully uninstalled llama-index-program-openai-0.1.7\n",
      "  Attempting uninstall: llama-index-question-gen-openai\n",
      "    Found existing installation: llama-index-question-gen-openai 0.1.3\n",
      "    Uninstalling llama-index-question-gen-openai-0.1.3:\n",
      "      Successfully uninstalled llama-index-question-gen-openai-0.1.3\n",
      "  Attempting uninstall: llama-index-multi-modal-llms-openai\n",
      "    Found existing installation: llama-index-multi-modal-llms-openai 0.1.8\n",
      "    Uninstalling llama-index-multi-modal-llms-openai-0.1.8:\n",
      "      Successfully uninstalled llama-index-multi-modal-llms-openai-0.1.8\n",
      "  Attempting uninstall: llama-index-cli\n",
      "    Found existing installation: llama-index-cli 0.1.13\n",
      "    Uninstalling llama-index-cli-0.1.13:\n",
      "      Successfully uninstalled llama-index-cli-0.1.13\n",
      "  Attempting uninstall: llama-index\n",
      "    Found existing installation: llama-index 0.10.57\n",
      "    Uninstalling llama-index-0.10.57:\n",
      "      Successfully uninstalled llama-index-0.10.57\n",
      "Successfully installed llama-cloud-0.0.17 llama-index-0.11.9 llama-index-agent-openai-0.3.1 llama-index-cli-0.3.1 llama-index-core-0.11.9 llama-index-embeddings-openai-0.2.4 llama-index-indices-managed-llama-cloud-0.3.0 llama-index-llms-openai-0.2.5 llama-index-multi-modal-llms-openai-0.2.0 llama-index-program-openai-0.2.0 llama-index-question-gen-openai-0.2.0 llama-index-readers-file-0.2.1 llama-index-readers-llama-parse-0.3.0 llama-parse-0.5.5 nltk-3.9.1\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython3.12 -m pip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip3.12 install -U llama-index --break-system-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2d697dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b8a55c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
    "\n",
    "from literalai import LiteralClient \n",
    "client = LiteralClient(api_key=\"\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "04270735",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.instrument_llamaindex()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "12291e8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-09-13 10:17:39--  https://github.com/user-attachments/files/16474262/data.zip\n",
      "Resolving github.com (github.com)... 140.82.121.3\n",
      "Connecting to github.com (github.com)|140.82.121.3|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://objects.githubusercontent.com/github-production-repository-file-5c1aeb/835341327/16474262?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240913%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240913T081739Z&X-Amz-Expires=300&X-Amz-Signature=a51ddb5e247535c64fe8105ca8a4370fec6eb67aac692ef73a89a2d388b034f7&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=835341327&response-content-disposition=attachment%3Bfilename%3Ddata.zip&response-content-type=application%2Fzip [following]\n",
      "--2024-09-13 10:17:39--  https://objects.githubusercontent.com/github-production-repository-file-5c1aeb/835341327/16474262?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240913%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240913T081739Z&X-Amz-Expires=300&X-Amz-Signature=a51ddb5e247535c64fe8105ca8a4370fec6eb67aac692ef73a89a2d388b034f7&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=835341327&response-content-disposition=attachment%3Bfilename%3Ddata.zip&response-content-type=application%2Fzip\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.110.133, 185.199.111.133, ...\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 94746 (93K) [application/zip]\n",
      "Saving to: ‘data.zip’\n",
      "\n",
      "data.zip            100%[===================>]  92,53K  --.-KB/s    in 0,009s  \n",
      "\n",
      "2024-09-13 10:17:40 (9,53 MB/s) - ‘data.zip’ saved [94746/94746]\n",
      "\n",
      "Archive:  data.zip\n",
      "  inflating: data/four_quadrants_of_conformism.txt  \n",
      "  inflating: data/haters.txt         \n",
      "  inflating: data/what_i_worked_on.txt  \n",
      "  inflating: data/two_kinds_of_moderate.txt  \n",
      "  inflating: data/the_best_essay.txt  \n",
      "  inflating: data/how_to_raise_money.txt  \n",
      "  inflating: data/investor_herd_dynamics.txt  \n",
      "  inflating: data/how_to_make_pgh_a_startup_hub.txt  \n",
      "  inflating: data/lies_we_tell_kids.txt  \n",
      "  inflating: data/being_a_noob.txt   \n",
      "  inflating: data/mean_people_fail.txt  \n"
     ]
    }
   ],
   "source": [
    "!wget https://github.com/user-attachments/files/16474262/data.zip -O data.zip\n",
    "!unzip -o data.zip\n",
    "!rm data.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "45804932",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict, Optional, Set, FrozenSet\n",
    "\n",
    "from llama_index.core.schema import BaseNode, TextNode\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "# constants\n",
    "DEFAULT_CHUNK_SIZE = 4096  # optionally splits documents into CHUNK_SIZE, then regroups them to demonstrate grouping algorithm\n",
    "DEFAULT_MAX_GROUP_SIZE = 20  # maximum number of documents in a group\n",
    "DEFAULT_SMALL_CHUNK_SIZE = 512  # small chunk size for generating embeddings\n",
    "DEFAULT_TOP_K = 8  # top k for retrieving\n",
    "\n",
    "\n",
    "def split_doc(chunk_size: int, documents: List[BaseNode]) -> List[TextNode]:\n",
    "    \"\"\"Splits documents into smaller pieces.\n",
    "\n",
    "    Args:\n",
    "        chunk_size (int): Chunk size\n",
    "        documents (List[BaseNode]): Documents\n",
    "\n",
    "    Returns:\n",
    "        List[TextNode]: Smaller chunks\n",
    "    \"\"\"\n",
    "    # split docs into tokens\n",
    "    text_parser = SentenceSplitter(chunk_size=chunk_size)\n",
    "    return text_parser.get_nodes_from_documents(documents)\n",
    "\n",
    "\n",
    "def group_docs(\n",
    "    nodes: List[str],\n",
    "    adj: Dict[str, List[str]],\n",
    "    max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
    ") -> Set[FrozenSet[str]]:\n",
    "    \"\"\"Groups documents.\n",
    "\n",
    "    Args:\n",
    "        nodes (List[str]): documents IDs\n",
    "        adj (Dict[str, List[str]]): related documents for each document; id -> list of doc strings\n",
    "        max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
    "    \"\"\"\n",
    "    docs = sorted(nodes, key=lambda node: len(adj[node]))\n",
    "    groups = set()  # set of set of IDs\n",
    "    for d in docs:\n",
    "        related_groups = set()\n",
    "        for r in adj[d]:\n",
    "            for g in groups:\n",
    "                if r in g:\n",
    "                    related_groups = related_groups.union(frozenset([g]))\n",
    "\n",
    "        gnew = {d}\n",
    "        related_groupsl = sorted(related_groups, key=lambda el: len(el))\n",
    "        for g in related_groupsl:\n",
    "            if max_group_size is None or len(gnew) + len(g) <= max_group_size:\n",
    "                gnew = gnew.union(g)\n",
    "                if g in groups:\n",
    "                    groups.remove(g)\n",
    "\n",
    "        groups.add(frozenset(gnew))\n",
    "\n",
    "    return groups\n",
    "\n",
    "\n",
    "def get_grouped_docs(\n",
    "    nodes: List[TextNode],\n",
    "    max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
    ") -> List[TextNode]:\n",
    "    \"\"\"Gets list of documents that are grouped.\n",
    "\n",
    "    Args:\n",
    "        nodes (t.List[TextNode]): Input list\n",
    "        max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
    "\n",
    "    Returns:\n",
    "        t.List[TextNode]: Output list\n",
    "    \"\"\"\n",
    "    # node IDs\n",
    "    nodes_str = [node.id_ for node in nodes]\n",
    "    # maps node ID -> related node IDs based on that node's relationships\n",
    "    adj: Dict[str, List[str]] = {\n",
    "        node.id_: [val.node_id for val in node.relationships.values()]\n",
    "        for node in nodes\n",
    "    }\n",
    "    # node ID -> node\n",
    "    nodes_dict = {node.id_: node for node in nodes}\n",
    "\n",
    "    res = group_docs(nodes_str, adj, max_group_size)\n",
    "\n",
    "    ret_nodes = []\n",
    "    for g in res:\n",
    "        cur_node = TextNode()\n",
    "\n",
    "        for node_id in g:\n",
    "            cur_node.text += nodes_dict[node_id].text + \"\\n\\n\"\n",
    "            cur_node.metadata.update(nodes_dict[node_id].metadata)\n",
    "\n",
    "        ret_nodes.append(cur_node)\n",
    "\n",
    "    return ret_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "62df709b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.retrievers import BaseRetriever\n",
    "from llama_index.core.vector_stores.simple import BasePydanticVectorStore\n",
    "from llama_index.core.schema import QueryBundle, NodeWithScore\n",
    "from llama_index.core.vector_stores.types import VectorStoreQuery\n",
    "from llama_index.core.settings import Settings\n",
    "\n",
    "\n",
    "class LongRAGRetriever(BaseRetriever):\n",
    "    \"\"\"Long RAG Retriever.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        grouped_nodes: List[TextNode],\n",
    "        small_toks: List[TextNode],\n",
    "        vector_store: BasePydanticVectorStore,\n",
    "        similarity_top_k: int = DEFAULT_TOP_K,\n",
    "    ) -> None:\n",
    "        \"\"\"Constructor.\n",
    "\n",
    "        Args:\n",
    "            grouped_nodes (List[TextNode]): Long retrieval units, nodes with docs grouped together based on relationships\n",
    "            small_toks (List[TextNode]): Smaller tokens\n",
    "            embed_model (BaseEmbedding, optional): Embed model. Defaults to None.\n",
    "            similarity_top_k (int, optional): Similarity top k. Defaults to 8.\n",
    "        \"\"\"\n",
    "        self._grouped_nodes = grouped_nodes\n",
    "        self._grouped_nodes_dict = {node.id_: node for node in grouped_nodes}\n",
    "        self._small_toks = small_toks\n",
    "        self._small_toks_dict = {node.id_: node for node in self._small_toks}\n",
    "\n",
    "        self._similarity_top_k = similarity_top_k\n",
    "        self._vec_store = vector_store\n",
    "        self._embed_model = Settings.embed_model\n",
    "\n",
    "    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:\n",
    "        \"\"\"Retrieves.\n",
    "\n",
    "        Args:\n",
    "            query_bundle (QueryBundle): query bundle\n",
    "\n",
    "        Returns:\n",
    "            List[NodeWithScore]: nodes with scores\n",
    "        \"\"\"\n",
    "        # make query\n",
    "        query_embedding = self._embed_model.get_query_embedding(\n",
    "            query_bundle.query_str\n",
    "        )\n",
    "        vector_store_query = VectorStoreQuery(\n",
    "            query_embedding=query_embedding, similarity_top_k=500\n",
    "        )\n",
    "\n",
    "        # query for answer\n",
    "        query_res = self._vec_store.query(vector_store_query)\n",
    "\n",
    "        # determine top parents of most similar children (these are long retrieval units)\n",
    "        top_parents_set: Set[str] = set()\n",
    "        top_parents: List[NodeWithScore] = []\n",
    "        for id_, similarity in zip(query_res.ids, query_res.similarities):\n",
    "            cur_node = self._small_toks_dict[id_]\n",
    "            parent_id = cur_node.ref_doc_id\n",
    "            if parent_id not in top_parents_set:\n",
    "                top_parents_set.add(parent_id)\n",
    "\n",
    "                parent_node = self._grouped_nodes_dict[parent_id]\n",
    "                node_with_score = NodeWithScore(\n",
    "                    node=parent_node, score=similarity\n",
    "                )\n",
    "                top_parents.append(node_with_score)\n",
    "\n",
    "                if len(top_parents_set) >= self._similarity_top_k:\n",
    "                    break\n",
    "\n",
    "        assert len(top_parents) == min(\n",
    "            self._similarity_top_k, len(self._grouped_nodes)\n",
    "        )\n",
    "\n",
    "        return top_parents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fd7ea8cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.llms import LLM\n",
    "from llama_index.core.workflow import Event\n",
    "\n",
    "\n",
    "class LoadNodeEvent(Event):\n",
    "    \"\"\"Event for loading nodes.\"\"\"\n",
    "\n",
    "    small_nodes: Iterable[TextNode]\n",
    "    grouped_nodes: list[TextNode]\n",
    "    index: VectorStoreIndex\n",
    "    similarity_top_k: int\n",
    "    llm: LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9b9c91f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Workflow,\n",
    "    step,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Context,\n",
    ")\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "\n",
    "class LongRAGWorkflow(Workflow):\n",
    "    \"\"\"Long RAG Workflow.\"\"\"\n",
    "\n",
    "    @step\n",
    "    async def ingest(self, ev: StartEvent) -> LoadNodeEvent | None:\n",
    "        \"\"\"Ingestion step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): Context\n",
    "            ev (StartEvent): start event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent | None: stop event with result\n",
    "        \"\"\"\n",
    "        data_dir: str = ev.get(\"data_dir\")\n",
    "        llm: LLM = ev.get(\"llm\")\n",
    "        chunk_size: int | None = ev.get(\"chunk_size\")\n",
    "        similarity_top_k: int = ev.get(\"similarity_top_k\")\n",
    "        small_chunk_size: int = ev.get(\"small_chunk_size\")\n",
    "        index: VectorStoreIndex | None = ev.get(\"index\")\n",
    "        index_kwargs: dict[str, t.Any] | None = ev.get(\"index_kwargs\")\n",
    "\n",
    "        if any(\n",
    "            i is None\n",
    "            for i in [data_dir, llm, similarity_top_k, small_chunk_size]\n",
    "        ):\n",
    "            return None\n",
    "\n",
    "        if not index:\n",
    "            docs = SimpleDirectoryReader(data_dir).load_data()\n",
    "            if chunk_size is not None:\n",
    "                nodes = split_doc(\n",
    "                    chunk_size, docs\n",
    "                )  # split documents into chunks of chunk_size\n",
    "                grouped_nodes = get_grouped_docs(\n",
    "                    nodes\n",
    "                )  # get list of nodes after grouping (groups are combined into one node), these are long retrieval units\n",
    "            else:\n",
    "                grouped_nodes = docs\n",
    "\n",
    "            # split large retrieval units into smaller nodes\n",
    "            small_nodes = split_doc(small_chunk_size, grouped_nodes)\n",
    "\n",
    "            index_kwargs = index_kwargs or {}\n",
    "            index = VectorStoreIndex(small_nodes, **index_kwargs)\n",
    "        else:\n",
    "            # get smaller nodes from index and form large retrieval units from these nodes\n",
    "            small_nodes = index.docstore.docs.values()\n",
    "            grouped_nodes = get_grouped_docs(small_nodes, None)\n",
    "\n",
    "        return LoadNodeEvent(\n",
    "            small_nodes=small_nodes,\n",
    "            grouped_nodes=grouped_nodes,\n",
    "            index=index,\n",
    "            similarity_top_k=similarity_top_k,\n",
    "            llm=llm,\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def make_query_engine(\n",
    "        self, ctx: Context, ev: LoadNodeEvent\n",
    "    ) -> StopEvent:\n",
    "        \"\"\"Query engine construction step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): context\n",
    "            ev (LoadNodeEvent): event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent: stop event\n",
    "        \"\"\"\n",
    "        # make retriever and query engine\n",
    "        retriever = LongRAGRetriever(\n",
    "            grouped_nodes=ev.grouped_nodes,\n",
    "            small_toks=ev.small_nodes,\n",
    "            similarity_top_k=ev.similarity_top_k,\n",
    "            vector_store=ev.index.vector_store,\n",
    "        )\n",
    "        query_eng = RetrieverQueryEngine.from_args(retriever, ev.llm)\n",
    "\n",
    "        return StopEvent(\n",
    "            result={\n",
    "                \"retriever\": retriever,\n",
    "                \"query_engine\": query_eng,\n",
    "                \"index\": ev.index,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def query(self, ctx: Context, ev: StartEvent) -> StopEvent | None:\n",
    "        \"\"\"Query step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): context\n",
    "            ev (StartEvent): start event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent | None: stop event with result\n",
    "        \"\"\"\n",
    "        query_str: str | None = ev.get(\"query_str\")\n",
    "        query_eng = ev.get(\"query_eng\")\n",
    "\n",
    "        if query_str is None:\n",
    "            return None\n",
    "\n",
    "        result = query_eng.query(query_str)\n",
    "        return StopEvent(result=result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "edcd273a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "wf = LongRAGWorkflow(timeout=60)\n",
    "llm = OpenAI(\"gpt-4o\")\n",
    "data_dir = \"data\"\n",
    "\n",
    "# initialize the workflow\n",
    "result = await wf.run(\n",
    "    data_dir=data_dir,\n",
    "    llm=llm,\n",
    "    chunk_size=DEFAULT_CHUNK_SIZE,\n",
    "    similarity_top_k=DEFAULT_TOP_K,\n",
    "    small_chunk_size=DEFAULT_SMALL_CHUNK_SIZE,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0b48445d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "To transform Pittsburgh into a startup hub, several strategies can be employed. Encouraging a youth-driven food boom, preserving historic buildings, and making the city more bicycle and pedestrian-friendly are key steps. Additionally, leveraging the presence of a first-rate research university like Carnegie Mellon University (CMU) can attract talent and foster innovation. The city should also focus on creating a supportive environment for startups by streamlining permit processes and promoting tolerance and diversity.\n",
       "\n",
       "Regarding the two types of moderates, they can be categorized as intentional and accidental moderates. Intentional moderates deliberately choose positions midway between the extremes of right and left. Accidental moderates, on the other hand, form their opinions independently on each issue, resulting in a distribution of views that average out to a moderate position."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "# run a query\n",
    "res = await wf.run(\n",
    "    query_str=\"How can Pittsburgh become a startup hub, and what are the two types of moderates?\",\n",
    "    query_eng=result[\"query_engine\"],\n",
    ")\n",
    "display(Markdown(str(res)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e2a30af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}