Introduction

Welcome to the second lesson of our "Scaling Up RAG with Vector Databases" course! In the previous unit, you explored how to break large documents into smaller chunks and attach useful metadata (like doc_id, chunk_id, and labels such as category). These chunks are essential for structuring data in a way that makes retrieval easier.

Now, let's build on that foundation by learning how to store these chunks in a vector database. One popular choice is ChromaDB—a specialized, open-source database designed for high-speed, semantic querying of vectors. By switching from keyword-based searches to semantic searches, your RAG system will retrieve relevant information more efficiently. Let’s start by understanding what vector databases are and why they matter for RAG.

Understanding Vector Databases

To work effectively with RAG systems, it's important to understand what a vector database is and how it works. A vector database stores data in the form of numerical vectors that capture the semantic essence of texts (or other data). The database then uses similarity metrics—rather than literal word matches—so that conceptually similar items are stored close together. This means searches on vector databases can retrieve contextually relevant results even when keywords are absent.

By leveraging approximate or exact nearest-neighbor strategies for similarity, vector databases can scale to handle millions or billions of vectors while still providing quick query responses. This makes them especially suitable for RAG systems, which rely on fast semantic lookups across large collections of text.

Now that you know what a vector database is, let’s see why it’s a crucial component for RAG pipelines.

Why We Need Vector Databases for RAG

Understanding the benefits of vector databases will help you appreciate their role in RAG systems. Here’s why they’re essential:

  1. Semantic Retrieval: By embedding text into vectors, queries can match documents based on meaning rather than strict keyword matches. This yields more accurate and context-sensitive search results.
  2. Scalability: Specialized vector databases handle large datasets efficiently, allowing you to store and query vast libraries of text chunks without sacrificing performance.
  3. Richer Context: Embeddings capture nuanced relationships among chunks, ensuring that related information is surfaced even when it doesn't use the exact same terms.
  4. Easy Updates: Vector databases (like ChromaDB) often allow you to add and remove chunks on the fly, so your collection stays in sync with new or evolving information.

With these advantages in mind, let’s move on to setting up ChromaDB and configuring it for your RAG workflow.

Setting Up ChromaDB and Basic Configuration

To start using ChromaDB, you need to initialize a client and create (or retrieve) a collection where your text chunks will be stored. The following code demonstrates how to do this in Rust, with comments explaining each step:

// Initialize a new ChromaDB client with default options (in-memory by default)
let client = ChromaClient::new(ChromaClientOptions::default()).await?;

// Get an existing collection or create a new one if it doesn't exist
let collection = client.get_or_create_collection(collection_name, None).await?;

// If there are no chunks to add, simply return the collection
if chunks.is_empty() {
    return Ok(collection);
}

This snippet sets up the connection to ChromaDB and ensures you have a collection ready to store your data. The ChromaClient handles communication with the database, while get_or_create_collection ensures you have a logical container for your text chunks and their embeddings. The early return for empty chunks prevents unnecessary operations.

Preparing Data and Adding Chunks to ChromaDB

Once your collection is ready, you need to prepare your data in the format ChromaDB expects. This involves extracting the text, generating unique IDs, attaching metadata, and creating embeddings. The following code walks through each step, with inline comments for clarity:

// Extract the text from each chunk for embedding and storage
let texts: Vec<String> = chunks.iter().map(|chunk| chunk.text.clone()).collect();

// Convert the texts to &str for compatibility with embedding and ChromaDB APIs
let documents: Vec<&str> = texts.iter().map(AsRef::as_ref).collect();

// Generate unique IDs for each chunk by combining doc_id and chunk_id
let ids_owned: Vec<String> = chunks.iter()
    .map(|chunk| format!("doc_{}_chunk_{}", chunk.doc_id, chunk.chunk_id))
    .collect();
let ids: Vec<&str> = ids_owned.iter().map(AsRef::as_ref).collect();

// Build metadata for each chunk (doc_id, chunk_id, category)
let metadatas = chunks.iter()
    .map(|chunk| {
        let mut map = serde_json::Map::new();
        map.insert("doc_id".to_string(), json!(chunk.doc_id));
        map.insert("chunk_id".to_string(), json!(chunk.chunk_id));
        map.insert("category".to_string(), chunk.category.clone().into());
        map
    })
    .collect();

// Generate vector embeddings for each document using your embedder
let embeddings = embedder.embed_texts(&documents)?;

// Bundle all data into a CollectionEntries struct for upsert
let entries = CollectionEntries {
    ids,
    embeddings: Some(embeddings),
    metadatas: Some(metadatas),
    documents: Some(documents),
};

// Upload all entries to the collection in a single batch operation
collection.upsert(entries, None).await?;

Here, you first extract the necessary fields from your chunk data. Each chunk is assigned a unique ID, and its metadata is structured for easy filtering and retrieval later. The embedder.embed_texts call transforms your text into vector representations, which are essential for semantic search. Finally, all prepared data is uploaded to ChromaDB in a single, efficient batch operation using upsert.

Managing Documents with Keyword-Based Deletion

As your collection grows, you may need to remove outdated or irrelevant documents. The following code demonstrates how to find and delete documents containing a specific keyword. Each part is explained with comments and additional context:

// Set up options to retrieve all documents, including only their text
let get_options = chromadb::collection::GetOptions {
    ids: Vec::new(), // Empty vector means fetch all documents
    where_metadata: None, // No metadata filter
    limit: None, // No limit on number of results
    offset: None, // No offset; start from the beginning
    where_document: None, // No document filter
    include: Some(vec!["documents".to_string()]), // Only fetch document texts
};

// Fetch all documents from the collection using the specified options
let data = collection.get(get_options).await?;

This snippet configures a query to retrieve all documents from your collection, but only includes the document text to minimize data transfer. This is useful when you want to scan for a keyword without loading unnecessary metadata or embeddings.

// Convert the keyword to lowercase for case-insensitive matching
let keyword_lower = keyword.to_lowercase();
let mut to_delete = Vec::new();

// Iterate over all documents and collect IDs of those containing the keyword
if let Some(documents) = &data.documents {
    for (i, doc_text) in documents.iter().enumerate() {
        if let Some(text) = doc_text {
            if text.to_lowercase().contains(&keyword_lower) {
                to_delete.push(data.ids[i].clone());
            }
        }
    }
}

Here, you loop through each document, checking if the text contains the target keyword (case-insensitive). If a match is found, the corresponding document ID is added to a list for deletion. This approach ensures you only target relevant documents for removal.

// If any matching documents were found, delete them in a single batch operation
if !to_delete.is_empty() {
    let to_delete_refs: Vec<&str> = to_delete.iter().map(|s| s.as_str()).collect();
    collection.delete(Some(to_delete_refs), None, None).await?;
}

Finally, if any documents matched the keyword, you delete them all at once using their IDs. This batch deletion is efficient and ensures your collection stays up to date without unnecessary overhead.

Conclusion and Next Steps

By storing text chunks in a vector database, you've laid the foundation for faster, more semantically aware retrieval. You know how to create, update, and manage a ChromaDB collection—crucial skills for any large-scale RAG system.

In the next lesson, you'll learn how to query the vector database to fetch the most relevant chunks and feed them into a language model. That's where the real magic of producing context-rich, accurate responses shines! For now, feel free to explore different embedding models or try adding and deleting a variety of chunks. When you're ready, proceed to the practice exercises to cement these concepts and further refine your RAG workflow.

Sign up
Join the 1M+ learners on CodeSignal
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal