Contextual Retrieval with VLLM and LlamaIndex

Suppose you build a RAG system. Sometimes the search returns chunks that look relevant but are not correct. The words match, but the context is wrong.

The cause is context loss during chunking. A chunk is a small piece of a document. When we split a document into chunks, each chunk loses its link to the whole document. The Contextual Retriever pattern adds this context back before we embed the chunks.

This post explains the idea in simple words. It also shows the technical steps. You can go to the section that is most useful to you.

The context loss problem

Imagine you find one page on the ground. It says: "The function returns embeddings from the last hidden layer." You cannot tell if it is about Python, JavaScript, machine learning, or database queries. You do not know which book the page came from.

Traditional RAG systems have the same problem. They cut documents into small pieces (chunks) so that they can search them. Each piece loses the context of its source. It is the same as when you tear pages out of books and put them in a pile.

An embedding is a list of numbers that shows the meaning of a text. Suppose you embed chunks without context. The vector space then becomes unclear. A chunk that says "The function returns embeddings" can match queries about any embedding function in any library. The embedding model captures the words, but not the semantic scope.

The result is low recall and irrelevant results. The correct chunk is in your database, but the query vector does not find it. Many similar but wrong chunks fill the embedding space.

The solution: contextual enrichment

The idea is simple. Before we save a chunk, we add a header that tells where the chunk came from. This is the same as writing the book title and chapter name at the top of each page. Anyone who finds the page then knows which book it is from.

We use VLLM, a fast engine that runs language models. VLLM reads each chunk and writes a short description of it. We store this description with the chunk. Searches then become more accurate.

The contextual enrichment pipeline has four stages. Each stage adds more semantic information to the chunk:

1. Document context: title, summary, and document type
2. Section context: headers and the subsection hierarchy
3. Neighbor context: the chunks around it, to keep the flow of the text
4. LLM-generated description: a short summary of the purpose of the chunk

Then we embed the enriched text. The vector now shows the content of the chunk. It also shows the position of the chunk in the document hierarchy.

Generate context with VLLM

Think of VLLM as a very fast librarian. For each page in your library, the librarian writes a short note. For example: "This page is from the Python documentation. It explains how the embed function works with custom models."

VLLM does this work very quickly. It can process thousands of chunks each minute. This makes it practical for large document collections.

VLLM gives the throughput that you need for context generation at production scale. Three settings are important: batch processing, model choice, and token limits.

from llama_index.llms.vllm import Vllm

# Configure VLLM endpoint for high throughput
llm = Vllm(
    api_url="<YOUR_API_URI>",
    model="meta-llama/Llama-2-7b-chat-hf",
    max_tokens=150  # Keep context descriptions concise
)

# The prompt template is critical
CONTEXT_PROMPT = f"""Given this document context:
Title: {doc.title}
Section: {section.header}

Describe what this chunk is about in one sentence:
{chunk.text}"""

context = llm.complete(CONTEXT_PROMPT).text

Store vectors with Milvus

Milvus is a vector database. Think of it as a filing cabinet. It does not sort files in alphabetical order. It sorts them by meaning. You search for "how to embed documents". Milvus finds all the pages about embedding, even if they use different words.

With contextual enrichment, each file also has a label that tells which book it came from. The cabinet knows the full context of each page. It gives you better results.

Milvus does vector similarity search. It also supports hybrid queries, which combine dense vectors and sparse keyword matching. The enriched context improves both.

from pymilvus import connections, Collection

# Connect to Milvus
connections.connect("default", host="localhost", port="19530")

# Create collection with contextual field
collection = Collection(
    "contextual_chunks",
    schema=CollectionSchema([
        FieldSchema("id", DataType.INT64, is_primary=True),
        FieldSchema("embedding", DataType.FLOAT_VECTOR, dim=768),
        FieldSchema("context", DataType.VARCHAR, max_length=500),
        FieldSchema("chunk_text", DataType.VARCHAR, max_length=2000),
        FieldSchema("doc_title", DataType.VARCHAR, max_length=200),
    ])
)

# Search with context-aware similarity
results = collection.search(
    data=[query_embedding],
    anns_field="embedding",
    param={"metric_type": "COSINE", "params": {"nprobe": 10}},
    limit=5,
    output_fields=["context", "chunk_text", "doc_title"]
)

The complete architecture

The diagram below shows how the parts connect. Document processing sends chunks to VLLM, which generates the context. The enriched chunks go to Milvus for vector storage. LlamaIndex manages the queries.

+----------------+     +----------------+     +----------------+
|   Documents    | --> |   Chunking     | --> |   VLLM         |
|   (PDF, MD)    |     |   (Semantic)   |     |   Context Gen  |
+----------------+     +----------------+     +----------------+
                                                     |
                                                     v
+----------------+     +----------------+     +----------------+
|   Query        | --> |   Similarity   | <-- |   Milvus       |
|   (User Input) |     |   Search       |     |   Vector Store |
+----------------+     +----------------+     +----------------+
                              |
                              v
                       +----------------+
                       |   LlamaIndex   |
                       |   Response     |
                       +----------------+

Summary

Remember these four points:

1. Context is important. Raw chunks lose the document hierarchy that gives them meaning. When you add context back, search relevance improves.
2. LLM-generated descriptions work. VLLM writes semantic descriptions at scale. Unclear chunks become well-labeled content.
3. Milvus works at production scale. Vector databases such as Milvus handle millions of vectors with sub-second latency. Real-time RAG needs this.
4. The enrichment pipeline is important. Document context, section headers, and neighbor chunks all improve the quality of the embeddings.

The full working notebook is on GitHub:Contextual-Retriever-working-notebook

I use this pattern in my RAG projects. It improves retrieval quality for different domains and document types. Try it and compare the results.