{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3058e9ca-07c3-4eef-b98c-bc2f2dbb9cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install -U langchain umap-learn scikit-learn langchain_community tiktoken langchain-openai langchainhub langchain-chroma langchain-anthropic"
   ]
  },
  {
   "attachments": {
    "72039e0c-e8c4-4b17-8780-04ad9fc584f3.png": {
     "image/png": "[encoded data, 161472 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "ea54c848-0df6-474e-b266-218a2acf67d3",
   "metadata": {},
   "source": [
    "# RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval\n",
    "\n",
    "The [RAPTOR](https://arxiv.org/pdf/2401.18059.pdf) paper presents an interesting approaching for indexing and retrieval of documents:\n",
    "\n",
    "* The `leafs` are a set of starting documents\n",
    "* Leafs are embedded and clustered\n",
    "* Clusters are then summarized into higher level (more abstract) consolidations of information across similar documents\n",
    "\n",
    "This process is done recursivly, resulting in a \"tree\" going from raw docs (`leafs`) to more abstract summaries.\n",
    " \n",
    "We can applying this at varying scales; `leafs` can be:\n",
    "\n",
    "* Text chunks from a single doc (as shown in the paper)\n",
    "* Full docs (as we show below)\n",
    "\n",
    "With longer context LLMs, it's possible to perform this over full documents. \n",
    "\n",
    "![Screenshot 2024-03-04 at 12.45.25 PM.png](attachment:72039e0c-e8c4-4b17-8780-04ad9fc584f3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "083dd961-b401-4fc6-867c-8f8950059b02",
   "metadata": {},
   "source": [
    "### Docs\n",
    "\n",
    "Let's apply this to LangChain's LCEL documentation.\n",
    "\n",
    "In this case, each `doc` is a unique web page of the LCEL docs.\n",
    "\n",
    "The context varies from < 2k tokens on up to > 10k tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b17c1331-373f-491d-8b53-ccf634e68c8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function matplotlib.pyplot.show(close=None, block=None)>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24164 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import tiktoken\n",
    "from bs4 import BeautifulSoup as Soup\n",
    "from langchain_community.document_loaders.recursive_url_loader import RecursiveUrlLoader\n",
    "\n",
    "\n",
    "def num_tokens_from_string(string: str, encoding_name: str) -> int:\n",
    "    \"\"\"Returns the number of tokens in a text string.\"\"\"\n",
    "    encoding = tiktoken.get_encoding(encoding_name)\n",
    "    num_tokens = len(encoding.encode(string))\n",
    "    return num_tokens\n",
    "\n",
    "\n",
    "# LCEL docs\n",
    "url = \"https://python.langchain.com/docs/expression_language/\"\n",
    "loader = RecursiveUrlLoader(\n",
    "    url=url, max_depth=20, extractor=lambda x: Soup(x, \"html.parser\").text\n",
    ")\n",
    "docs = loader.load()\n",
    "\n",
    "# LCEL w/ PydanticOutputParser (outside the primary LCEL docs)\n",
    "url = \"https://python.langchain.com/docs/modules/model_io/output_parsers/quick_start\"\n",
    "loader = RecursiveUrlLoader(\n",
    "    url=url, max_depth=1, extractor=lambda x: Soup(x, \"html.parser\").text\n",
    ")\n",
    "docs_pydantic = loader.load()\n",
    "\n",
    "# LCEL w/ Self Query (outside the primary LCEL docs)\n",
    "url = \"https://python.langchain.com/docs/modules/data_connection/retrievers/self_query/\"\n",
    "loader = RecursiveUrlLoader(\n",
    "    url=url, max_depth=1, extractor=lambda x: Soup(x, \"html.parser\").text\n",
    ")\n",
    "docs_sq = loader.load()\n",
    "\n",
    "# Doc texts\n",
    "docs.extend([*docs_pydantic, *docs_sq])\n",
    "docs_texts = [d.page_content for d in docs]\n",
    "\n",
    "# Calculate the number of tokens for each document\n",
    "counts = [num_tokens_from_string(d, \"cl100k_base\") for d in docs_texts]\n",
    "\n",
    "# Plotting the histogram of token counts\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.hist(counts, bins=30, color=\"blue\", edgecolor=\"black\", alpha=0.7)\n",
    "plt.title(\"Histogram of Token Counts\")\n",
    "plt.xlabel(\"Token Count\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.grid(axis=\"y\", alpha=0.75)\n",
    "\n",
    "# Display the histogram\n",
    "plt.show"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "70750603-ec82-4439-9b32-d22014b5ff2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num tokens in all context: 68705\n"
     ]
    }
   ],
   "source": [
    "# Doc texts concat\n",
    "d_sorted = sorted(docs, key=lambda x: x.metadata[\"source\"])\n",
    "d_reversed = list(reversed(d_sorted))\n",
    "concatenated_content = \"\\n\\n\\n --- \\n\\n\\n\".join(\n",
    "    [doc.page_content for doc in d_reversed]\n",
    ")\n",
    "print(\n",
    "    \"Num tokens in all context: %s\"\n",
    "    % num_tokens_from_string(concatenated_content, \"cl100k_base\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "25ca3cf2-0f6b-40f9-a2ff-285a8dcb33dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Doc texts split\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "chunk_size_tok = 2000\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=chunk_size_tok, chunk_overlap=0\n",
    ")\n",
    "texts_split = text_splitter.split_text(concatenated_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "797a5469-0942-45a5-adb6-f12e05d76798",
   "metadata": {},
   "source": [
    "## Models\n",
    "\n",
    "We can test various models, including the new [Claude3](https://www.anthropic.com/news/claude-3-family) family.\n",
    "\n",
    "Be sure to set the relevant API keys:\n",
    "\n",
    "* `ANTHROPIC_API_KEY`\n",
    "* `OPENAI_API_KEY`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "033e71d3-5dc8-42a3-a0b7-4df116048c14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "embd = OpenAIEmbeddings()\n",
    "\n",
    "# from langchain_openai import ChatOpenAI\n",
    "\n",
    "# model = ChatOpenAI(temperature=0, model=\"gpt-4-1106-preview\")\n",
    "\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "\n",
    "model = ChatAnthropic(temperature=0, model=\"claude-3-opus-20240229\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c63db01-cf95-4c17-ae5d-8dc7267ad58a",
   "metadata": {},
   "source": [
    "### Tree Constrution\n",
    "\n",
    "The clustering approach in tree construction includes a few interesting ideas.\n",
    "\n",
    "**GMM (Gaussian Mixture Model)** \n",
    "\n",
    "- Model the distribution of data points across different clusters\n",
    "- Optimal number of clusters by evaluating the model's Bayesian Information Criterion (BIC)\n",
    "\n",
    "**UMAP (Uniform Manifold Approximation and Projection)** \n",
    "\n",
    "- Supports clustering\n",
    "- Reduces the dimensionality of high-dimensional data\n",
    "- UMAP helps to highlight the natural grouping of data points based on their similarities\n",
    "\n",
    "**Local and Global Clustering** \n",
    "\n",
    "- Used to analyze data at different scales\n",
    "- Both fine-grained and broader patterns within the data are captured effectively\n",
    "\n",
    "**Thresholding** \n",
    "\n",
    "- Apply in the context of GMM to determine cluster membership\n",
    "- Based on the probability distribution (assignment of data points to ≥ 1 cluster)\n",
    "---\n",
    "\n",
    "Code for GMM and thresholding is from Sarthi et al, as noted in the below two sources:\n",
    " \n",
    "* [Origional repo](https://github.com/parthsarthi03/raptor/blob/master/raptor/cluster_tree_builder.py)\n",
    "* [Minor tweaks](https://github.com/run-llama/llama_index/blob/main/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/clustering.py)\n",
    "\n",
    "Full credit to both authors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a849980c-27d4-48e0-87a0-c2a5143cb8c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, List, Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import umap\n",
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from sklearn.mixture import GaussianMixture\n",
    "\n",
    "RANDOM_SEED = 224  # Fixed seed for reproducibility\n",
    "\n",
    "### --- Code from citations referenced above (added comments and docstrings) --- ###\n",
    "\n",
    "\n",
    "def global_cluster_embeddings(\n",
    "    embeddings: np.ndarray,\n",
    "    dim: int,\n",
    "    n_neighbors: Optional[int] = None,\n",
    "    metric: str = \"cosine\",\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Perform global dimensionality reduction on the embeddings using UMAP.\n",
    "\n",
    "    Parameters:\n",
    "    - embeddings: The input embeddings as a numpy array.\n",
    "    - dim: The target dimensionality for the reduced space.\n",
    "    - n_neighbors: Optional; the number of neighbors to consider for each point.\n",
    "                   If not provided, it defaults to the square root of the number of embeddings.\n",
    "    - metric: The distance metric to use for UMAP.\n",
    "\n",
    "    Returns:\n",
    "    - A numpy array of the embeddings reduced to the specified dimensionality.\n",
    "    \"\"\"\n",
    "    if n_neighbors is None:\n",
    "        n_neighbors = int((len(embeddings) - 1) ** 0.5)\n",
    "    return umap.UMAP(\n",
    "        n_neighbors=n_neighbors, n_components=dim, metric=metric\n",
    "    ).fit_transform(embeddings)\n",
    "\n",
    "\n",
    "def local_cluster_embeddings(\n",
    "    embeddings: np.ndarray, dim: int, num_neighbors: int = 10, metric: str = \"cosine\"\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Perform local dimensionality reduction on the embeddings using UMAP, typically after global clustering.\n",
    "\n",
    "    Parameters:\n",
    "    - embeddings: The input embeddings as a numpy array.\n",
    "    - dim: The target dimensionality for the reduced space.\n",
    "    - num_neighbors: The number of neighbors to consider for each point.\n",
    "    - metric: The distance metric to use for UMAP.\n",
    "\n",
    "    Returns:\n",
    "    - A numpy array of the embeddings reduced to the specified dimensionality.\n",
    "    \"\"\"\n",
    "    return umap.UMAP(\n",
    "        n_neighbors=num_neighbors, n_components=dim, metric=metric\n",
    "    ).fit_transform(embeddings)\n",
    "\n",
    "\n",
    "def get_optimal_clusters(\n",
    "    embeddings: np.ndarray, max_clusters: int = 50, random_state: int = RANDOM_SEED\n",
    ") -> int:\n",
    "    \"\"\"\n",
    "    Determine the optimal number of clusters using the Bayesian Information Criterion (BIC) with a Gaussian Mixture Model.\n",
    "\n",
    "    Parameters:\n",
    "    - embeddings: The input embeddings as a numpy array.\n",
    "    - max_clusters: The maximum number of clusters to consider.\n",
    "    - random_state: Seed for reproducibility.\n",
    "\n",
    "    Returns:\n",
    "    - An integer representing the optimal number of clusters found.\n",
    "    \"\"\"\n",
    "    max_clusters = min(max_clusters, len(embeddings))\n",
    "    n_clusters = np.arange(1, max_clusters)\n",
    "    bics = []\n",
    "    for n in n_clusters:\n",
    "        gm = GaussianMixture(n_components=n, random_state=random_state)\n",
    "        gm.fit(embeddings)\n",
    "        bics.append(gm.bic(embeddings))\n",
    "    return n_clusters[np.argmin(bics)]\n",
    "\n",
    "\n",
    "def GMM_cluster(embeddings: np.ndarray, threshold: float, random_state: int = 0):\n",
    "    \"\"\"\n",
    "    Cluster embeddings using a Gaussian Mixture Model (GMM) based on a probability threshold.\n",
    "\n",
    "    Parameters:\n",
    "    - embeddings: The input embeddings as a numpy array.\n",
    "    - threshold: The probability threshold for assigning an embedding to a cluster.\n",
    "    - random_state: Seed for reproducibility.\n",
    "\n",
    "    Returns:\n",
    "    - A tuple containing the cluster labels and the number of clusters determined.\n",
    "    \"\"\"\n",
    "    n_clusters = get_optimal_clusters(embeddings)\n",
    "    gm = GaussianMixture(n_components=n_clusters, random_state=random_state)\n",
    "    gm.fit(embeddings)\n",
    "    probs = gm.predict_proba(embeddings)\n",
    "    labels = [np.where(prob > threshold)[0] for prob in probs]\n",
    "    return labels, n_clusters\n",
    "\n",
    "\n",
    "def perform_clustering(\n",
    "    embeddings: np.ndarray,\n",
    "    dim: int,\n",
    "    threshold: float,\n",
    ") -> List[np.ndarray]:\n",
    "    \"\"\"\n",
    "    Perform clustering on the embeddings by first reducing their dimensionality globally, then clustering\n",
    "    using a Gaussian Mixture Model, and finally performing local clustering within each global cluster.\n",
    "\n",
    "    Parameters:\n",
    "    - embeddings: The input embeddings as a numpy array.\n",
    "    - dim: The target dimensionality for UMAP reduction.\n",
    "    - threshold: The probability threshold for assigning an embedding to a cluster in GMM.\n",
    "\n",
    "    Returns:\n",
    "    - A list of numpy arrays, where each array contains the cluster IDs for each embedding.\n",
    "    \"\"\"\n",
    "    if len(embeddings) <= dim + 1:\n",
    "        # Avoid clustering when there's insufficient data\n",
    "        return [np.array([0]) for _ in range(len(embeddings))]\n",
    "\n",
    "    # Global dimensionality reduction\n",
    "    reduced_embeddings_global = global_cluster_embeddings(embeddings, dim)\n",
    "    # Global clustering\n",
    "    global_clusters, n_global_clusters = GMM_cluster(\n",
    "        reduced_embeddings_global, threshold\n",
    "    )\n",
    "\n",
    "    all_local_clusters = [np.array([]) for _ in range(len(embeddings))]\n",
    "    total_clusters = 0\n",
    "\n",
    "    # Iterate through each global cluster to perform local clustering\n",
    "    for i in range(n_global_clusters):\n",
    "        # Extract embeddings belonging to the current global cluster\n",
    "        global_cluster_embeddings_ = embeddings[\n",
    "            np.array([i in gc for gc in global_clusters])\n",
    "        ]\n",
    "\n",
    "        if len(global_cluster_embeddings_) == 0:\n",
    "            continue\n",
    "        if len(global_cluster_embeddings_) <= dim + 1:\n",
    "            # Handle small clusters with direct assignment\n",
    "            local_clusters = [np.array([0]) for _ in global_cluster_embeddings_]\n",
    "            n_local_clusters = 1\n",
    "        else:\n",
    "            # Local dimensionality reduction and clustering\n",
    "            reduced_embeddings_local = local_cluster_embeddings(\n",
    "                global_cluster_embeddings_, dim\n",
    "            )\n",
    "            local_clusters, n_local_clusters = GMM_cluster(\n",
    "                reduced_embeddings_local, threshold\n",
    "            )\n",
    "\n",
    "        # Assign local cluster IDs, adjusting for total clusters already processed\n",
    "        for j in range(n_local_clusters):\n",
    "            local_cluster_embeddings_ = global_cluster_embeddings_[\n",
    "                np.array([j in lc for lc in local_clusters])\n",
    "            ]\n",
    "            indices = np.where(\n",
    "                (embeddings == local_cluster_embeddings_[:, None]).all(-1)\n",
    "            )[1]\n",
    "            for idx in indices:\n",
    "                all_local_clusters[idx] = np.append(\n",
    "                    all_local_clusters[idx], j + total_clusters\n",
    "                )\n",
    "\n",
    "        total_clusters += n_local_clusters\n",
    "\n",
    "    return all_local_clusters\n",
    "\n",
    "\n",
    "### --- Our code below --- ###\n",
    "\n",
    "\n",
    "def embed(texts):\n",
    "    \"\"\"\n",
    "    Generate embeddings for a list of text documents.\n",
    "\n",
    "    This function assumes the existence of an `embd` object with a method `embed_documents`\n",
    "    that takes a list of texts and returns their embeddings.\n",
    "\n",
    "    Parameters:\n",
    "    - texts: List[str], a list of text documents to be embedded.\n",
    "\n",
    "    Returns:\n",
    "    - numpy.ndarray: An array of embeddings for the given text documents.\n",
    "    \"\"\"\n",
    "    text_embeddings = embd.embed_documents(texts)\n",
    "    text_embeddings_np = np.array(text_embeddings)\n",
    "    return text_embeddings_np\n",
    "\n",
    "\n",
    "def embed_cluster_texts(texts):\n",
    "    \"\"\"\n",
    "    Embeds a list of texts and clusters them, returning a DataFrame with texts, their embeddings, and cluster labels.\n",
    "\n",
    "    This function combines embedding generation and clustering into a single step. It assumes the existence\n",
    "    of a previously defined `perform_clustering` function that performs clustering on the embeddings.\n",
    "\n",
    "    Parameters:\n",
    "    - texts: List[str], a list of text documents to be processed.\n",
    "\n",
    "    Returns:\n",
    "    - pandas.DataFrame: A DataFrame containing the original texts, their embeddings, and the assigned cluster labels.\n",
    "    \"\"\"\n",
    "    text_embeddings_np = embed(texts)  # Generate embeddings\n",
    "    cluster_labels = perform_clustering(\n",
    "        text_embeddings_np, 10, 0.1\n",
    "    )  # Perform clustering on the embeddings\n",
    "    df = pd.DataFrame()  # Initialize a DataFrame to store the results\n",
    "    df[\"text\"] = texts  # Store original texts\n",
    "    df[\"embd\"] = list(text_embeddings_np)  # Store embeddings as a list in the DataFrame\n",
    "    df[\"cluster\"] = cluster_labels  # Store cluster labels\n",
    "    return df\n",
    "\n",
    "\n",
    "def fmt_txt(df: pd.DataFrame) -> str:\n",
    "    \"\"\"\n",
    "    Formats the text documents in a DataFrame into a single string.\n",
    "\n",
    "    Parameters:\n",
    "    - df: DataFrame containing the 'text' column with text documents to format.\n",
    "\n",
    "    Returns:\n",
    "    - A single string where all text documents are joined by a specific delimiter.\n",
    "    \"\"\"\n",
    "    unique_txt = df[\"text\"].tolist()\n",
    "    return \"--- --- \\n --- --- \".join(unique_txt)\n",
    "\n",
    "\n",
    "def embed_cluster_summarize_texts(\n",
    "    texts: List[str], level: int\n",
    ") -> Tuple[pd.DataFrame, pd.DataFrame]:\n",
    "    \"\"\"\n",
    "    Embeds, clusters, and summarizes a list of texts. This function first generates embeddings for the texts,\n",
    "    clusters them based on similarity, expands the cluster assignments for easier processing, and then summarizes\n",
    "    the content within each cluster.\n",
    "\n",
    "    Parameters:\n",
    "    - texts: A list of text documents to be processed.\n",
    "    - level: An integer parameter that could define the depth or detail of processing.\n",
    "\n",
    "    Returns:\n",
    "    - Tuple containing two DataFrames:\n",
    "      1. The first DataFrame (`df_clusters`) includes the original texts, their embeddings, and cluster assignments.\n",
    "      2. The second DataFrame (`df_summary`) contains summaries for each cluster, the specified level of detail,\n",
    "         and the cluster identifiers.\n",
    "    \"\"\"\n",
    "\n",
    "    # Embed and cluster the texts, resulting in a DataFrame with 'text', 'embd', and 'cluster' columns\n",
    "    df_clusters = embed_cluster_texts(texts)\n",
    "\n",
    "    # Prepare to expand the DataFrame for easier manipulation of clusters\n",
    "    expanded_list = []\n",
    "\n",
    "    # Expand DataFrame entries to document-cluster pairings for straightforward processing\n",
    "    for index, row in df_clusters.iterrows():\n",
    "        for cluster in row[\"cluster\"]:\n",
    "            expanded_list.append(\n",
    "                {\"text\": row[\"text\"], \"embd\": row[\"embd\"], \"cluster\": cluster}\n",
    "            )\n",
    "\n",
    "    # Create a new DataFrame from the expanded list\n",
    "    expanded_df = pd.DataFrame(expanded_list)\n",
    "\n",
    "    # Retrieve unique cluster identifiers for processing\n",
    "    all_clusters = expanded_df[\"cluster\"].unique()\n",
    "\n",
    "    print(f\"--Generated {len(all_clusters)} clusters--\")\n",
    "\n",
    "    # Summarization\n",
    "    template = \"\"\"Here is a sub-set of LangChain Expression Language doc. \n",
    "    \n",
    "    LangChain Expression Language provides a way to compose chain in LangChain.\n",
    "    \n",
    "    Give a detailed summary of the documentation provided.\n",
    "    \n",
    "    Documentation:\n",
    "    {context}\n",
    "    \"\"\"\n",
    "    prompt = ChatPromptTemplate.from_template(template)\n",
    "    chain = prompt | model | StrOutputParser()\n",
    "\n",
    "    # Format text within each cluster for summarization\n",
    "    summaries = []\n",
    "    for i in all_clusters:\n",
    "        df_cluster = expanded_df[expanded_df[\"cluster\"] == i]\n",
    "        formatted_txt = fmt_txt(df_cluster)\n",
    "        summaries.append(chain.invoke({\"context\": formatted_txt}))\n",
    "\n",
    "    # Create a DataFrame to store summaries with their corresponding cluster and level\n",
    "    df_summary = pd.DataFrame(\n",
    "        {\n",
    "            \"summaries\": summaries,\n",
    "            \"level\": [level] * len(summaries),\n",
    "            \"cluster\": list(all_clusters),\n",
    "        }\n",
    "    )\n",
    "\n",
    "    return df_clusters, df_summary\n",
    "\n",
    "\n",
    "def recursive_embed_cluster_summarize(\n",
    "    texts: List[str], level: int = 1, n_levels: int = 3\n",
    ") -> Dict[int, Tuple[pd.DataFrame, pd.DataFrame]]:\n",
    "    \"\"\"\n",
    "    Recursively embeds, clusters, and summarizes texts up to a specified level or until\n",
    "    the number of unique clusters becomes 1, storing the results at each level.\n",
    "\n",
    "    Parameters:\n",
    "    - texts: List[str], texts to be processed.\n",
    "    - level: int, current recursion level (starts at 1).\n",
    "    - n_levels: int, maximum depth of recursion.\n",
    "\n",
    "    Returns:\n",
    "    - Dict[int, Tuple[pd.DataFrame, pd.DataFrame]], a dictionary where keys are the recursion\n",
    "      levels and values are tuples containing the clusters DataFrame and summaries DataFrame at that level.\n",
    "    \"\"\"\n",
    "    results = {}  # Dictionary to store results at each level\n",
    "\n",
    "    # Perform embedding, clustering, and summarization for the current level\n",
    "    df_clusters, df_summary = embed_cluster_summarize_texts(texts, level)\n",
    "\n",
    "    # Store the results of the current level\n",
    "    results[level] = (df_clusters, df_summary)\n",
    "\n",
    "    # Determine if further recursion is possible and meaningful\n",
    "    unique_clusters = df_summary[\"cluster\"].nunique()\n",
    "    if level < n_levels and unique_clusters > 1:\n",
    "        # Use summaries as the input texts for the next level of recursion\n",
    "        new_texts = df_summary[\"summaries\"].tolist()\n",
    "        next_level_results = recursive_embed_cluster_summarize(\n",
    "            new_texts, level + 1, n_levels\n",
    "        )\n",
    "\n",
    "        # Merge the results from the next level into the current results dictionary\n",
    "        results.update(next_level_results)\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f0d8cd3e-cd49-484d-9617-1b9811cc08b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--Generated 7 clusters--\n",
      "--Generated 1 clusters--\n"
     ]
    }
   ],
   "source": [
    "# Build tree\n",
    "leaf_texts = docs_texts\n",
    "results = recursive_embed_cluster_summarize(leaf_texts, level=1, n_levels=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e80d7098-5d16-4fa6-837c-968e5c9f118d",
   "metadata": {},
   "source": [
    "The paper reports best performance from `collapsed tree retrieval`. \n",
    "\n",
    "This involves flattening the tree structure into a single layer and then applying a k-nearest neighbors (kNN) search across all nodes simultaneously. \n",
    "\n",
    "We do simply do this below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d28ba9e6-9124-41a8-b4fd-55a6ef4ac062",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "\n",
    "# Initialize all_texts with leaf_texts\n",
    "all_texts = leaf_texts.copy()\n",
    "\n",
    "# Iterate through the results to extract summaries from each level and add them to all_texts\n",
    "for level in sorted(results.keys()):\n",
    "    # Extract summaries from the current level's DataFrame\n",
    "    summaries = results[level][1][\"summaries\"].tolist()\n",
    "    # Extend all_texts with the summaries from the current level\n",
    "    all_texts.extend(summaries)\n",
    "\n",
    "# Now, use all_texts to build the vectorstore with Chroma\n",
    "vectorstore = Chroma.from_texts(texts=all_texts, embedding=embd)\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d497627-44c6-41f7-bb63-1d858d3f188f",
   "metadata": {},
   "source": [
    "Now we can using our flattened, indexed tree in a RAG chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9d6c894b-b3a3-4a01-b779-3e98ea382ff5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Here is a code example of how to define a RAG (Retrieval Augmented Generation) chain in LangChain:\\n\\n```python\\nfrom langchain.vectorstores import FAISS\\nfrom langchain.embeddings import OpenAIEmbeddings\\nfrom langchain.prompts import ChatPromptTemplate\\nfrom langchain.chat_models import ChatOpenAI\\nfrom langchain.output_parsers import StrOutputParser\\n\\n# Load documents into vector store\\nvectorstore = FAISS.from_texts(\\n    [\"harrison worked at kensho\"], embedding=OpenAIEmbeddings()\\n)\\nretriever = vectorstore.as_retriever()\\n\\n# Define prompt template\\ntemplate = \"\"\"Answer the question based only on the following context:\\n{context}\\nQuestion: {question}\"\"\"\\nprompt = ChatPromptTemplate.from_template(template)\\n\\n# Define model and output parser\\nmodel = ChatOpenAI()\\noutput_parser = StrOutputParser()\\n\\n# Define RAG chain\\nchain = (\\n    {\"context\": retriever, \"question\": RunnablePassthrough()}\\n    | prompt\\n    | model '"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain import hub\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "# Prompt\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "\n",
    "# Post-processing\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "\n",
    "# Chain\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | model\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# Question\n",
    "rag_chain.invoke(\"How to define a RAG chain? Give me a specific code example.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c585b37-ad83-4069-8f5d-4a6a3e15128d",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/1dabf475-1675-4494-b16c-928fbf079851/r"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
