Retrieval for LLMs

The Retrieval Mechanism

At its core, the retrieval step is a semantic search process. It's the engine that finds the needles of relevant information in the haystack of your document collection. The process unfolds in a precise sequence:

1. Receive Query: The system takes the user's question (e.g., "What are Python's naming conventions?") as a plain text string.
2. Embed Query: The text query is passed through the same embedding model that was used to index the document chunks. This converts the query into a vector that exists in the same high-dimensional "meaning space" as your documents. This is a crucial step; using different models would be like trying to compare coordinates from two different maps.
3. Similarity Search: This query vector is sent to the vector store. The store then performs a similarity search (e.g., using cosine similarity) to find the vectors in its index that are "closest" to the query vector.
4. Return Top-K Chunks: The search doesn't return everything—it returns the top-k most similar document chunks, where k is a number you can configure. These chunks represent the information most likely to contain the answer to the user's query.

The following diagram perfectly illustrates this flow from a query to a search within the vector space.

To understand the mechanics of this search, the video below offers an excellent explanation of how vector databases enable this powerful semantic search capability.

From Theory to Code: A Simple Implementation

Before we use a high-level framework, it's useful to understand the retrieval logic in its most basic form. As a software engineer, you can appreciate seeing the fundamental algorithm. A simplified retrieval process is just a search and sort operation.

The following article contains a toy example that demonstrates this beautifully.

This simple example, sorted(index.items(), key=lambda x: abs(x[1]-q_vec))[:top_k], is the heart of retrieval. Of course, real-world implementations use more sophisticated distance metrics (like cosine similarity) and highly optimized data structures (like HNSW), but the core idea is the same: find the indexed data points that are closest to the query.

Implementing Retrieval with LlamaIndex

Now, let's implement this using a robust framework. We will use LlamaIndex, a popular library designed specifically for building RAG applications. It provides high-level abstractions that handle the complexities of embedding, indexing, and retrieving.

For this exercise, we will assume you have already completed the "Ingestion" or "Indexing" process from our architectural diagram. This means you have a directory on your disk (e.g., ./storage) containing a persisted vector store index, which you would have created using the techniques from Module 4.

Our task is to load this index, create a dedicated retriever object, and use it to fetch relevant documents for a given query.

The process involves three main steps:

1. Load the Index: If the index is already persisted, we load it into memory. If not, we create it for the first time.
2. Create a Retriever: From the loaded index, we create a retriever object. This object is specialized for one task: fetching documents based on a query. This is distinct from a full query_engine, which we'll see later, as the engine also handles passing the context to the LLM and generating a final answer.
3. Retrieve Documents: We call the retriever's retrieve() method with our query string.

The following article provides all the necessary code, which we will adapt slightly to focus purely on the retrieval step.

Our Retrieval Script

Drawing from the RealPython guide, here is a Python script that isolates and performs the retrieval step. Note the small but important change: we use index.as_retriever() and retriever.retrieve() to stop after fetching the documents, before they are sent to an LLM.

from pathlib import Path
from llama_index.core import (
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
import os




# --- Configuration ---
# You will need to set your OpenAI API key for this to work.
# os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"

PERSIST_DIR = Path("./storage")
DATA_FILE = Path("./data/pep8.rst") # Using a sample data file




# --- Index Loading/Creation ---
def get_index(persist_dir: Path, data_file: Path):
    """Load the index from storage or create and persist it."""
    if persist_dir.exists():
        storage_context = StorageContext.from_defaults(
            persist_dir=str(persist_dir),
        )
        index = load_index_from_storage(storage_context)
        print(f"Index loaded from {persist_dir}")
    else:



        # Create the directory if it doesn't exist
        persist_dir.mkdir(parents=True, exist_ok=True)



        # Load data
        reader = SimpleDirectoryReader(input_files=[str(data_file)])
        documents = reader.load_data()



        # Create and persist index
        index = VectorStoreIndex.from_documents(documents)
        index.storage_context.persist(persist_dir=str(persist_dir))
        print(f"Index created and persisted to {persist_dir}")
    return index




# --- Main Retrieval Logic ---
def main():



    # Get the index
    index = get_index(persist_dir=PERSIST_DIR, data_file=DATA_FILE)




    # Create a retriever object from the index
    # We can configure it to return the top 3 most similar chunks
    retriever = index.as_retriever(similarity_top_k=3)




    # Define a query
    query = "Summarize the naming conventions in Python."




    # Retrieve relevant document chunks
    retrieved_nodes = retriever.retrieve(query)

    print(f"\n--- Query ---\n{query}\n")
    print(f"--- Retrieved {len(retrieved_nodes)} nodes ---")




    # Print the retrieved nodes and their similarity scores
    for i, node_with_score in enumerate(retrieved_nodes):
        print(f"\n--- Node {i+1} (Score: {node_with_score.score:.4f}) ---")
        print(node_with_score.get_content())

if __name__ == "__main__":
    main()

When you run this script (after setting up a sample data directory and your API key), the output will be the raw text chunks that the vector store determined were most relevant to your query, along with their similarity scores. You have successfully implemented the retrieval step.

To see this same process explained in a video format and using a different framework (LangChain with ChromaDB), the following walkthrough is highly recommended. It reinforces that the underlying principles are universal.

Conclusion

In this lesson, you have moved from a high-level architectural diagram to a concrete, working implementation of the retrieval step. You've seen how a user's text query is transformed and used to perform a semantic search against an indexed knowledge base, returning the most relevant pieces of information.

Key Takeaways:

• Retrieval is a 4-step process: Receive query, embed query, perform similarity search, and return the top-k results.
• The same embedding model must be used for both indexing your documents and embedding the user's query at runtime.
• Frameworks like LlamaIndex provide powerful abstractions like the retriever object, which simplifies the implementation of this step.
• The output of the retrieval step is a ranked list of document chunks (the context), which are not yet a final answer.

You have successfully fetched the raw materials needed to answer a question. In our next lesson, we will tackle the "A" and "G" in RAG: you will learn how to construct an augmented prompt that combines this retrieved context with the user's query and then pass it to an LLM to generate a grounded, final response.