raptor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval\n",
    "\n",
    "This notebook shows how to use an implementation of RAPTOR with llama-index, leveraging the RAPTOR llama-pack.\n",
    "\n",
    "RAPTOR works by recursively clustering and summarizing clusters in layers for retrieval.\n",
    "\n",
    "There two retrieval modes:\n",
    "- tree_traversal -- traversing the tree of clusters, performing top-k at each level in the tree.\n",
    "- collapsed -- treat the entire tree as a giant pile of nodes, perform simple top-k.\n",
    "\n",
    "See [the paper](https://arxiv.org/abs/2401.18059) for full algorithm details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: llama-index in /home/harish/.local/lib/python3.10/site-packages (0.10.9)\n",
      "Requirement already satisfied: llama-index-packs-raptor in /home/harish/.local/lib/python3.10/site-packages (0.1.3)\n",
      "Requirement already satisfied: llama-index-vector-stores-qdrant in /home/harish/.local/lib/python3.10/site-packages (0.2.8)\n",
      "Requirement already satisfied: chromadb in /home/harish/.local/lib/python3.10/site-packages (0.4.24)\n",
      "Collecting llama-index-vector-stores-chroma\n",
      "  Downloading llama_index_vector_stores_chroma-0.1.8-py3-none-any.whl (4.8 kB)\n",
      "Requirement already satisfied: llama-index-multi-modal-llms-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.4)\n",
      "Requirement already satisfied: llama-index-indices-managed-llama-cloud<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.4)\n",
      "Requirement already satisfied: llama-index-readers-llama-parse<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.3)\n",
      "Requirement already satisfied: llama-index-llms-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.12)\n",
      "Requirement already satisfied: llama-index-question-gen-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.3)\n",
      "Requirement already satisfied: llama-index-program-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.4)\n",
      "Requirement already satisfied: llama-index-readers-file<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.1)\n",
      "Requirement already satisfied: llama-index-embeddings-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.6)\n",
      "Requirement already satisfied: llama-index-core<0.11.0,>=0.10.8.post1 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.10.20.post2)\n",
      "Requirement already satisfied: llama-index-agent-openai<0.2.0,>=0.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.1.5)\n",
      "Requirement already satisfied: llama-index-legacy<0.10.0,>=0.9.48 in /home/harish/.local/lib/python3.10/site-packages (from llama-index) (0.9.48)\n",
      "Requirement already satisfied: umap-learn>=0.5.5 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-packs-raptor) (0.5.6)\n",
      "Requirement already satisfied: scikit-learn in /home/harish/.local/lib/python3.10/site-packages (from llama-index-packs-raptor) (1.4.0)\n",
      "Requirement already satisfied: qdrant-client<2.0.0,>=1.7.1 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-vector-stores-qdrant) (1.9.1)\n",
      "Requirement already satisfied: grpcio<2.0.0,>=1.60.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-vector-stores-qdrant) (1.64.0)\n",
      "Requirement already satisfied: fastapi>=0.95.2 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.103.0)\n",
      "Requirement already satisfied: mmh3>=4.0.1 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (4.1.0)\n",
      "Requirement already satisfied: tokenizers>=0.13.2 in /usr/local/lib/python3.10/dist-packages (from chromadb) (0.15.1)\n",
      "Requirement already satisfied: onnxruntime>=1.14.1 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.17.1)\n",
      "Requirement already satisfied: posthog>=2.4.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (3.5.0)\n",
      "Requirement already satisfied: chroma-hnswlib==0.7.3 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.7.3)\n",
      "Requirement already satisfied: pulsar-client>=3.1.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (3.4.0)\n",
      "Requirement already satisfied: build>=1.0.3 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.2.1)\n",
      "Requirement already satisfied: overrides>=7.3.1 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (7.7.0)\n",
      "Requirement already satisfied: orjson>=3.9.12 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (3.9.14)\n",
      "Requirement already satisfied: bcrypt>=4.0.1 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (4.1.2)\n",
      "Requirement already satisfied: opentelemetry-sdk>=1.2.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.22.0)\n",
      "Requirement already satisfied: kubernetes>=28.1.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (29.0.0)\n",
      "Requirement already satisfied: tenacity>=8.2.3 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (8.2.3)\n",
      "Requirement already satisfied: uvicorn[standard]>=0.18.3 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.27.0.post1)\n",
      "Requirement already satisfied: numpy>=1.22.5 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.26.3)\n",
      "Requirement already satisfied: importlib-resources in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (6.3.0)\n",
      "Requirement already satisfied: typing-extensions>=4.5.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (4.9.0)\n",
      "Requirement already satisfied: pydantic>=1.9 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (2.6.4)\n",
      "Requirement already satisfied: requests>=2.28 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (2.31.0)\n",
      "Requirement already satisfied: pypika>=0.48.9 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.48.9)\n",
      "Requirement already satisfied: tqdm>=4.65.0 in /usr/local/lib/python3.10/dist-packages (from chromadb) (4.66.1)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-grpc>=1.2.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.22.0)\n",
      "Requirement already satisfied: PyYAML>=6.0.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (6.0.1)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-fastapi>=0.41b0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.43b0)\n",
      "Requirement already satisfied: opentelemetry-api>=1.2.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (1.22.0)\n",
      "Requirement already satisfied: typer>=0.9.0 in /home/harish/.local/lib/python3.10/site-packages (from chromadb) (0.9.4)\n",
      "Requirement already satisfied: pyproject_hooks in /home/harish/.local/lib/python3.10/site-packages (from build>=1.0.3->chromadb) (1.0.0)\n",
      "Requirement already satisfied: packaging>=19.1 in /home/harish/.local/lib/python3.10/site-packages (from build>=1.0.3->chromadb) (23.2)\n",
      "Requirement already satisfied: tomli>=1.1.0 in /home/harish/.local/lib/python3.10/site-packages (from build>=1.0.3->chromadb) (2.0.1)\n",
      "Requirement already satisfied: starlette<0.28.0,>=0.27.0 in /home/harish/.local/lib/python3.10/site-packages (from fastapi>=0.95.2->chromadb) (0.27.0)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (1.7.0)\n",
      "Requirement already satisfied: six>=1.9.0 in /usr/lib/python3/dist-packages (from kubernetes>=28.1.0->chromadb) (1.16.0)\n",
      "Requirement already satisfied: requests-oauthlib in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (1.3.1)\n",
      "Requirement already satisfied: oauthlib>=3.2.2 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (3.2.2)\n",
      "Requirement already satisfied: urllib3>=1.24.2 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (1.26.18)\n",
      "Requirement already satisfied: python-dateutil>=2.5.3 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (2.8.2)\n",
      "Requirement already satisfied: google-auth>=1.0.1 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (2.27.0)\n",
      "Requirement already satisfied: certifi>=14.05.14 in /home/harish/.local/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb) (2024.2.2)\n",
      "Requirement already satisfied: openai>=1.1.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.23.6)\n",
      "Requirement already satisfied: nest-asyncio<2.0.0,>=1.5.8 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.6.0)\n",
      "Requirement already satisfied: pandas in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (2.2.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (2023.12.2)\n",
      "Requirement already satisfied: llamaindex-py-client<0.2.0,>=0.1.13 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.1.13)\n",
      "Requirement already satisfied: tiktoken>=0.3.3 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.5.2)\n",
      "Requirement already satisfied: deprecated>=1.2.9.3 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.2.14)\n",
      "Requirement already satisfied: SQLAlchemy[asyncio]>=1.4.49 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (2.0.27)\n",
      "Requirement already satisfied: dirtyjson<2.0.0,>=1.0.8 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.0.8)\n",
      "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (3.8.1)\n",
      "Requirement already satisfied: dataclasses-json in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.6.4)\n",
      "Requirement already satisfied: networkx>=3.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (3.2.1)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.6 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (3.9.3)\n",
      "Requirement already satisfied: pillow>=9.0.0 in /usr/lib/python3/dist-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (9.0.1)\n",
      "Requirement already satisfied: httpx in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.24.1)\n",
      "Requirement already satisfied: typing-inspect>=0.8.0 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.9.0)\n",
      "Requirement already satisfied: bs4<0.0.3,>=0.0.2 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-readers-file<0.2.0,>=0.1.0->llama-index) (0.0.2)\n",
      "Requirement already satisfied: beautifulsoup4<5.0.0,>=4.12.3 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-readers-file<0.2.0,>=0.1.0->llama-index) (4.12.3)\n",
      "Requirement already satisfied: pymupdf<2.0.0,>=1.23.21 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-readers-file<0.2.0,>=0.1.0->llama-index) (1.23.26)\n",
      "Requirement already satisfied: llama-parse<0.4.0,>=0.3.3 in /home/harish/.local/lib/python3.10/site-packages (from llama-index-readers-llama-parse<0.2.0,>=0.1.0->llama-index) (0.3.9)\n",
      "Requirement already satisfied: coloredlogs in /home/harish/.local/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in /home/harish/.local/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb) (24.3.25)\n",
      "Requirement already satisfied: protobuf in /home/harish/.local/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb) (4.25.3)\n",
      "Requirement already satisfied: sympy in /home/harish/.local/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb) (1.12)\n",
      "Requirement already satisfied: importlib-metadata<7.0,>=6.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-api>=1.2.0->chromadb) (6.11.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-common==1.22.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.22.0)\n",
      "Requirement already satisfied: opentelemetry-proto==1.22.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.22.0)\n",
      "Requirement already satisfied: googleapis-common-protos~=1.52 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.62.0)\n",
      "Requirement already satisfied: backoff<3.0.0,>=1.10.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (2.2.1)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-asgi==0.43b0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.43b0)\n",
      "Requirement already satisfied: opentelemetry-semantic-conventions==0.43b0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.43b0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation==0.43b0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.43b0)\n",
      "Requirement already satisfied: opentelemetry-util-http==0.43b0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.43b0)\n",
      "Requirement already satisfied: wrapt<2.0.0,>=1.0.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation==0.43b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (1.16.0)\n",
      "Requirement already satisfied: setuptools>=16.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation==0.43b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (68.2.2)\n",
      "Requirement already satisfied: asgiref~=3.0 in /home/harish/.local/lib/python3.10/site-packages (from opentelemetry-instrumentation-asgi==0.43b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (3.8.1)\n",
      "Requirement already satisfied: monotonic>=1.5 in /usr/lib/python3/dist-packages (from posthog>=2.4.0->chromadb) (1.6)\n",
      "Requirement already satisfied: pydantic-core==2.16.3 in /home/harish/.local/lib/python3.10/site-packages (from pydantic>=1.9->chromadb) (2.16.3)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /home/harish/.local/lib/python3.10/site-packages (from pydantic>=1.9->chromadb) (0.6.0)\n",
      "Requirement already satisfied: grpcio-tools>=1.41.0 in /home/harish/.local/lib/python3.10/site-packages (from qdrant-client<2.0.0,>=1.7.1->llama-index-vector-stores-qdrant) (1.62.2)\n",
      "Requirement already satisfied: portalocker<3.0.0,>=2.7.0 in /home/harish/.local/lib/python3.10/site-packages (from qdrant-client<2.0.0,>=1.7.1->llama-index-vector-stores-qdrant) (2.8.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests>=2.28->chromadb) (3.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/harish/.local/lib/python3.10/site-packages (from requests>=2.28->chromadb) (3.3.2)\n",
      "Requirement already satisfied: huggingface_hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from tokenizers>=0.13.2->chromadb) (0.20.3)\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in /home/harish/.local/lib/python3.10/site-packages (from typer>=0.9.0->chromadb) (8.1.7)\n",
      "Requirement already satisfied: pynndescent>=0.5 in /home/harish/.local/lib/python3.10/site-packages (from umap-learn>=0.5.5->llama-index-packs-raptor) (0.5.12)\n",
      "Requirement already satisfied: numba>=0.51.2 in /home/harish/.local/lib/python3.10/site-packages (from umap-learn>=0.5.5->llama-index-packs-raptor) (0.59.1)\n",
      "Requirement already satisfied: scipy>=1.3.1 in /home/harish/.local/lib/python3.10/site-packages (from umap-learn>=0.5.5->llama-index-packs-raptor) (1.12.0)\n",
      "Requirement already satisfied: joblib>=1.2.0 in /home/harish/.local/lib/python3.10/site-packages (from scikit-learn->llama-index-packs-raptor) (1.3.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /home/harish/.local/lib/python3.10/site-packages (from scikit-learn->llama-index-packs-raptor) (3.2.0)\n",
      "Requirement already satisfied: h11>=0.8 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.14.0)\n",
      "Requirement already satisfied: python-dotenv>=0.13 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (1.0.0)\n",
      "Requirement already satisfied: watchfiles>=0.13 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.21.0)\n",
      "Requirement already satisfied: websockets>=10.4 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (12.0)\n",
      "Requirement already satisfied: httptools>=0.5.0 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.6.1)\n",
      "Requirement already satisfied: uvloop!=0.15.0,!=0.15.1,>=0.14.0 in /home/harish/.local/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.19.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (6.0.5)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.3.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (4.0.3)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.4.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/harish/.local/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.6->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (23.2.0)\n",
      "Requirement already satisfied: soupsieve>1.2 in /home/harish/.local/lib/python3.10/site-packages (from beautifulsoup4<5.0.0,>=4.12.3->llama-index-readers-file<0.2.0,>=0.1.0->llama-index) (2.5)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/harish/.local/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.3.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /home/harish/.local/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (4.9)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /home/harish/.local/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (5.3.2)\n",
      "Requirement already satisfied: httpcore<0.18.0,>=0.15.0 in /home/harish/.local/lib/python3.10/site-packages (from httpx->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (0.17.3)\n",
      "Requirement already satisfied: sniffio in /home/harish/.local/lib/python3.10/site-packages (from httpx->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.3.0)\n",
      "Requirement already satisfied: h2<5,>=3 in /home/harish/.local/lib/python3.10/site-packages (from httpx->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (4.1.0)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface_hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb) (3.13.1)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/lib/python3/dist-packages (from importlib-metadata<7.0,>=6.0->opentelemetry-api>=1.2.0->chromadb) (1.0.0)\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-core<0.11.0,>=0.10.8.post1->llama-index) (2023.12.25)\n",
      "Requirement already satisfied: llvmlite<0.43,>=0.42.0dev0 in /home/harish/.local/lib/python3.10/site-packages (from numba>=0.51.2->umap-learn>=0.5.5->llama-index-packs-raptor) (0.42.0)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /usr/lib/python3/dist-packages (from openai>=1.1.0->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.7.0)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /home/harish/.local/lib/python3.10/site-packages (from openai>=1.1.0->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (4.2.0)\n",
      "Requirement already satisfied: PyMuPDFb==1.23.22 in /home/harish/.local/lib/python3.10/site-packages (from pymupdf<2.0.0,>=1.23.21->llama-index-readers-file<0.2.0,>=0.1.0->llama-index) (1.23.22)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /home/harish/.local/lib/python3.10/site-packages (from SQLAlchemy[asyncio]>=1.4.49->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (3.0.3)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/harish/.local/lib/python3.10/site-packages (from typing-inspect>=0.8.0->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.0.0)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /home/harish/.local/lib/python3.10/site-packages (from coloredlogs->onnxruntime>=1.14.1->chromadb) (10.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /home/harish/.local/lib/python3.10/site-packages (from dataclasses-json->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (3.20.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/harish/.local/lib/python3.10/site-packages (from pandas->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (2023.4)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /home/harish/.local/lib/python3.10/site-packages (from pandas->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (2023.4)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/harish/.local/lib/python3.10/site-packages (from sympy->onnxruntime>=1.14.1->chromadb) (1.3.0)\n",
      "Requirement already satisfied: exceptiongroup>=1.0.2 in /home/harish/.local/lib/python3.10/site-packages (from anyio<5,>=3.5.0->openai>=1.1.0->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (1.2.0)\n",
      "Requirement already satisfied: hpack<5,>=4.0 in /home/harish/.local/lib/python3.10/site-packages (from h2<5,>=3->httpx->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (4.0.0)\n",
      "Requirement already satisfied: hyperframe<7,>=6.0 in /home/harish/.local/lib/python3.10/site-packages (from h2<5,>=3->httpx->llama-index-core<0.11.0,>=0.10.8.post1->llama-index) (6.0.1)\n",
      "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /home/harish/.local/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.5.1)\n",
      "Installing collected packages: llama-index-vector-stores-chroma\n",
      "Successfully installed llama-index-vector-stores-chroma-0.1.8\n"
     ]
    }
   ],
   "source": [
    "!pip install llama-index llama-index-packs-raptor llama-index-vector-stores-qdrant chromadb  llama-index-vector-stores-chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.packs.raptor import RaptorPack\n",
    "\n",
    "# optionally download the pack to inspect/modify it yourself!\n",
    "# from llama_index.core.llama_pack import download_llama_pack\n",
    "# RaptorPack = download_llama_pack(\"RaptorPack\", \"./raptor_pack\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-cHgnNwpihYLM884UqVkDT3BlbkFJGw0AgW8etJSuz6s13Le1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing the Clusters/Hierarchy Tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(input_files=[\"./Form Master Services Agreement (Outsourcing).DOCX\"]).load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating embeddings for level 0.\n",
      "Performing clustering for level 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/harish/.local/lib/python3.10/site-packages/numba/np/ufunc/parallel.py:371: NumbaWarning: \u001b[1mThe TBB threading layer requires TBB version 2021 update 6 or later i.e., TBB_INTERFACE_VERSION >= 12060. Found TBB_INTERFACE_VERSION = 12050. The TBB threading layer is disabled.\u001b[0m\n",
      "  warnings.warn(problem)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating summaries for level 0 with 24 clusters.\n",
      "Level 0 created summaries/clusters: 24\n",
      "Generating embeddings for level 1.\n",
      "Performing clustering for level 1.\n",
      "Generating summaries for level 1 with 4 clusters.\n",
      "Level 1 created summaries/clusters: 4\n",
      "Generating embeddings for level 2.\n",
      "Performing clustering for level 2.\n",
      "Generating summaries for level 2 with 1 clusters.\n",
      "Level 2 created summaries/clusters: 1\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore # type: ignore\n",
    "import chromadb\n",
    "\n",
    "client = chromadb.PersistentClient(path=\"./raptor_paper_db\")\n",
    "collection = client.get_or_create_collection(\"raptor\")\n",
    "\n",
    "vector_store = ChromaVectorStore(chroma_collection=collection)\n",
    "\n",
    "raptor_pack = RaptorPack(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(\n",
    "        model=\"text-embedding-3-small\"\n",
    "    ),  # used for embedding clusters\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1),  # used for generating summaries\n",
    "    vector_store=vector_store,  # used for storage\n",
    "    similarity_top_k=2,  # top k for each layer, or overall top-k for collapsed\n",
    "    mode=\"collapsed\",  # sets default mode\n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=400, chunk_overlap=50)\n",
    "    ],  # transformations applied for ingestion\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "The agreement outlines the terms related to changes in charges, intellectual property rights, ownership of work products, disclosure of inventions, and software provisions between the Provider and the Client. It specifies that any changes in charges must be agreed upon and adjusted accordingly. The agreement also addresses the ownership of software, work products, and modifications, with the Client retaining all rights and interests. It further discusses the use of Provider Proprietary Materials and the disclosure of inventions made during the agreement. Additionally, it covers the installation of upgrades and modifications to software, as well as the granting of rights and licenses related to Provider Patents within specific industries.\n"
     ]
    }
   ],
   "source": [
    "nodes = raptor_pack.run(\"What are the intellectual property rights of the vendor?\", mode=\"collapsed\")\n",
    "print(len(nodes))\n",
    "print(nodes[0].text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieved parent IDs from level 2: ['0f0c386b-12f5-460c-9fb2-5d65b0024c7c']\n",
      "Retrieved 2 from parents at level 2.\n",
      "Retrieved parent IDs from level 1: ['750e2166-637d-4a6a-8b6c-86b69ba6ad97', 'd6ead990-a89c-420d-b1bb-8108516f4b8f']\n",
      "Retrieved 4 from parents at level 1.\n",
      "Retrieved parent IDs from level 0: ['a9cc1f53-70ae-4e05-bfb5-e409ccd914f6', '637bef95-1443-4125-8ae0-4a7793896c0a']\n",
      "Retrieved 4 from parents at level 0.\n",
      "4\n",
      "The territorial extent of the rights in the Work Product assigned to [XXXXXXXX] by Provider and/or the Provider Personnel under this Agreement shall extend to all the countries in the world.  The assignment of the Intellectual Property Rights in the Work Product by Provider and/or the Provider Personnel to [XXXXXXXX] shall be royalty-free absolute, irrevocable and perpetual.  \n",
      "\n",
      "With respect to any Services performed in India, the Parties agree that, without limitation of any other [XXXXXXXX] rights or remedies under the Agreement, the following provisions shall apply: (i) Provider acknowledges that notwithstanding the provisions of Section 19(4) of the Indian Copyright Act, 1957, any assignment to [XXXXXXXX] of all Intellectual Property Rights to Work Product in connection with this Agreement, any SOW and/or Work Order shall not lapse nor the right transferred therein revert to Provider and/or the Provider Personnel even if [XXXXXXXX] and the successors and permitted assigns of [XXXXXXXX] do not exercise the rights under assignment within a period of one (1) year from the date of such assignment.  Provider shall not, and shall ensure that its Affiliates and the Provider Personnel do not, take any steps against [XXXXXXXX] and the successors and permitted assigns of [XXXXXXXX] under Section 19A of the Copyright Act, 1957.  For the avoidance of doubt, the Parties agree that in the event of the commencement of any winding up proceeding by or against Provider under the Indian Companies Act, 1956, [XXXXXXXX] shall be entitled to retain all of its rights under this Agreement, any SOW and/or Work Order. \n",
      "\n",
      "Provider Software.\n"
     ]
    }
   ],
   "source": [
    "nodes = raptor_pack.run(\n",
    "    \"What are the intellectual property rights of the vendor?\", mode=\"tree_traversal\"\n",
    ")\n",
    "print(len(nodes))\n",
    "print(nodes[0].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading\n",
    "\n",
    "Since we saved to a vector store, we can also use it again! (For local vector stores, there is a `persist` and `from_persist_dir` method on the retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.packs.raptor import RaptorRetriever\n",
    "\n",
    "retriever = RaptorRetriever(\n",
    "    [],\n",
    "    embed_model=OpenAIEmbedding(\n",
    "        model=\"text-embedding-3-small\"\n",
    "    ),  # used for embedding clusters\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1),  # used for generating summaries\n",
    "    vector_store=vector_store,  # used for storage\n",
    "    similarity_top_k=2,  # top k for each layer, or overall top-k for collapsed\n",
    "    mode=\"tree_traversal\",  # sets default mode\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if using a default vector store\n",
    "# retriever.persist(\"./persist\")\n",
    "# retriever = RaptorRetriever.from_persist_dir(\"./persist\", ...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "query_engine = RetrieverQueryEngine.from_args(\n",
    "    retriever, llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"Explain how termination works in the agreement?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Termination in the agreement can occur for various reasons such as expiration, termination for convenience, termination for cause, termination due to breach, termination for damages cap exceeded, or termination for non-payment. In the event of termination, certain obligations and liabilities of both parties may continue, and termination may lead to the automatic termination of all work orders then in effect. Additionally, there are provisions for termination assistance to facilitate the orderly transfer of services back to the provider or its designee, with charges for such assistance specified in the agreement.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-4aB9_5sa-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}