tools/rag/llama_index_context_window/llama_index_context_window.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Context Enrichment Window for Document Retrieval with Llama Index\n",
        "\n",
        "This notebook introduces a new RAG technique that leverages a context enrichment window to improve document retrieval. By incorporating surrounding context into text chunks, the system aims to return more coherent and contextually rich results when querying the vector index."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## System Introduction\n",
        "\n",
        "The system integrates a specialized context enrichment mechanism into the Llama Index pipeline. It is designed to ingest documents from a corpus, enrich each document chunk with a configurable window of adjacent sentences, and build a vector index. In retrieval, it uses custom postprocessing to reconstruct the most coherent piece of text around the query match. This integration not only retains the advantages of vector search but also addresses its common issue of returning isolated text fragments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Underlying Concept\n",
        "\n",
        "Traditional vector-based retrieval may return segments of text that lack context. This technique enriches document chunks by adding a \"window\" of surrounding sentences. The idea is to provide additional context to each chunk, ensuring that when a document is retrieved, it comes with enough contextual information to be coherent and complete. This approach also allows for flexible adjustment of the window size to suit different datasets and application needs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
       "## System Components\n",
       "\n",
       "1. **Document Retrieval Module:** Loads a persisted llama_index vector store and retrieves documents relevant to a query. \n",
       "   - *Example Prompt for Retrieval:* \"Find the latest market trends using the llama_index tool.\"\n",
       "\n",
       "2. **Corpus Ingestion Module:** Reads documents from a specified directory using the SimpleDirectoryReader, then creates or updates the llama_index vector store.\n",
       "   - **Bulk Ingestion:** You can ingest a corpus by directly uploading files to the `/tools/rag/llama_index_context_window/corpus` folder. To start the batch processing, simply make an API call:\n",
       "     ```bash\n",
       "     curl -X POST \\\n",
       "       http://localhost:5000/llama-index-ingest-corpus \\\n",
       "       -H \"Content-Type: application/json\" \\\n",
       "       -d '{\"isContextWindow\": true}'\n",
       "     ```\n",
       "     The script processes the information in the files, transferring it to the vector and graph database. By default, the SimpleDirectoryReader attempts to read any files it encounters (treating them as text), and explicitly supports file types such as:\n",
       "     - .csv (Comma-Separated Values)\n",
       "     - .docx (Microsoft Word)\n",
       "     - .epub (EPUB eBook format)\n",
       "     - .hwp (Hangul Word Processor)\n",
       "     - .ipynb (Jupyter Notebook)\n",
       "     - .jpeg, .jpg (JPEG image)\n",
       "     - .mbox (MBOX email archive)\n",
       "     - .md (Markdown)\n",
       "     - .mp3, .mp4 (Audio and video)\n",
       "     - .pdf (Portable Document Format)\n",
       "     - .png (Portable Network Graphics)\n",
       "     - .ppt, .pptm, .pptx (Microsoft PowerPoint)\n",
       "     \n",
       "     For more details, refer to the [SimpleDirectoryReader Documentation](https://docs.llamaindex.ai/en/stable/module_guides/loading/simpledirectoryreader/)."
      ]
     },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## How It Works\n",
        "\n",
        "### 1. Document Ingestion\n",
        "\n",
        "- The ingestion module reads documents from a corpus directory using the `SimpleDirectoryReader`.\n",
        "- A `SentenceWindowNodeParser` is applied via an `IngestionPipeline` to enrich each document chunk with a window of adjacent sentences. This provides additional context around each text segment.\n",
        "- The module then either creates a new index or updates an existing one by inserting the enriched nodes, and persists the index to disk.\n",
        "\n",
        "### 2. Document Retrieval\n",
        "\n",
        "- The retrieval module loads the persisted index from disk.\n",
        "- A custom metadata replacement postprocessor scans the enriched context window to extract the segment of text surrounding the original sentence, based on a configurable maximum number of adjacent characters.\n",
        "- If a similarity cutoff is specified, an additional postprocessor refines the results to ensure only highly relevant nodes are returned.\n",
        "- The final output is a list of processed text fragments that maintain coherent context for downstream applications."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Workflow Diagram\n",
        "\n",
        "![LLama Index Context Window Workflow](./llama_index_context_window_workflow.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## System Advantages\n",
        "\n",
        "- **Contextual Richness:** Enriching document chunks with a surrounding sentence window improves the coherence of the retrieved content.\n",
        "- **Enhanced Retrieval Quality:** By providing extra context, the technique mitigates the problem of isolated text fragments, leading to more accurate responses.\n",
        "- **Flexibility:** The context window size is configurable, allowing the system to be tuned for various types of datasets and use cases.\n",
        "- **Seamless Integration:** It builds upon the existing Llama Index framework, maintaining the benefits of vector search while addressing its limitations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Practical Benefits\n",
        "\n",
        "- **More Coherent Results:** Retrieval returns enriched, context-aware text fragments that are easier to understand and more useful for tasks like question answering.\n",
        "- **Mitigated Fragmentation:** The approach overcomes the common issue of isolated text fragments in vector-based retrieval systems.\n",
        "- **Adjustable Context:** Users can adjust the window size based on the complexity and structure of their documents, optimizing retrieval quality.\n",
        "\n",
        "Overall, this technique enhances the performance of document retrieval systems by ensuring that the output retains necessary contextual information, leading to improved downstream processing and interpretation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Implementation Insights\n",
        "\n",
        "- The **retrieve.py** module leverages llama_index's `StorageContext` and `load_index_from_storage` along with a custom postprocessor (`CustomMetadataReplacementPostProcessor`) to process and trim enriched text chunks during document retrieval.\n",
        "- The **ingest_corpus.py** module uses `SimpleDirectoryReader` and `SentenceWindowNodeParser` via an `IngestionPipeline` to ingest documents with context window enrichment and update or create a vector index."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Parameters\n",
        "\n",
        "**LLAMA_INDEX_CONTEXT_WINDOW_TOP_K_RAG_RETRIEVE:**\n",
        "This environment variable determines the number of top enriched document chunks retrieved during a query in the context window retrieval system. It sets the `similarity_top_k` parameter for the query engine, balancing the breadth and precision of retrieval.\n",
        "\n",
        "**LLAMA_INDEX_CONTEXT_WINDOW_MAX_ADJACENT_CHARS_RAG_RETRIEVE:**\n",
        "This variable specifies the maximum number of characters to include before and after the original sentence when reconstructing the context from the enriched window. It ensures that the returned text fragment maintains sufficient surrounding context for coherence.\n",
        "\n",
        "**LLAMA_INDEX_CONTEXT_WINDOW_SIZE_INGEST:**\n",
        "This parameter defines the number of adjacent sentences (on each side) to include during the ingestion process. It is used by the `SentenceWindowNodeParser` to enrich document chunks with additional context, thereby improving the quality of retrieval."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Conclusion\n",
        "\n",
        "Integrating a dynamic context enrichment window into the Llama Index workflow revolutionizes the way document retrieval systems operate. This approach adeptly stitches together neighboring sentences to create a more complete and coherent narrative around query matches, thereby overcoming the fragmentation common in conventional vector search. By fusing robust ingestion, refined vector indexing, and custom postprocessing, the system delivers retrieval outputs that capture the full essence of the source material. Ultimately, this innovative method offers unparalleled flexibility and precision, empowering users to tackle complex data challenges with confidence."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}