Building CSV RAG with Rig and Rust 🔥🔥🔥

Introduction

In this article, we're going to look at how you can use the rig Rust AI framework to create an application that is able to load a CSV file, embed it into a vector store and have an LLM answer questions based on provided context from our vector store. We'll serve the application in the form of a command-line tool.

Interested in trying it out for yourself? Check out the Github repo.

What is RAG?

Retrieval-Augmented Generation (RAG) is a technique that combines the power of information retrieval with the generative capabilities of large language models (LLMs). The process typically involves two main components: retrieving relevant knowledge from a pre-existing database or knowledge base and generating contextually accurate answers based on that information. When a question or query is posed, the model first retrieves relevant snippets or pieces of information from the knowledge base. It then augments its generative model with this newly acquired knowledge, enabling it to provide a more accurate, informed response.

In contrast to traditional LLMs that rely purely on their trained parameters, RAG systems dynamically incorporate external data, ensuring that responses are grounded in factual and up-to-date knowledge. When a question is asked, the relevant knowledge from a pre-existing dataset is retrieved and fed into the model. The LLM then generates an answer that combines the retrieved information and its generative abilities, aiming to offer a detailed and accurate response. This hybrid approach allows RAG models to provide answers that are not only linguistically fluent but also well-supported by factual data.

Getting started

To get started, make sure you have the Rust programming language installed.

Next, you'll need to create a new project using cargo init:

cargo init csv-rag
cd csv-rag

Let's get building!

Before we start, let's add the relevant dependencies which we need:

cargo add serde rig-core tokio csv -F \
serde/derive,tokio/macros,tokio/rt-multi-thread

This adds the following dependencies:

• rig-core - The rig library.
• tokio - an asynchronous Rust runtime. We additionally add the macros and rt-multi-thread feature as we want to use the macro.
• csv - The Rust crate for using CSVs.

Parsing CSVs

Before we do anything else, we'll need to declare our struct type as well as derive the correct macros for it. Note that the embedding trait for Rig also depends on Eq (which also depends on PartialEq), hence the derivation.

#[derive(serde::Deserialize, serde::Serialize, Debug, Eq, PartialEq)]
struct Record {
    first_name: String,
    last_name: String,
    email: String,
    role: String,
    salary: u32, //salary in dollars
}

Next to be able to use Rig's embeddings API, we need to implement the Embed trait. This trait simply defines a method for what to embed. Because we don't have a singular field we want to embed and will need the whole record, we can simply just implement fmt::Display for the struct then format the whole struct into a string (which we can then embed).

impl std::fmt::Display for Record {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self {
            first_name,
            last_name,
            email,
            role,
            salary,
        } = self;
        write!(
            f,
            "First name: {first_name}\nLast name: {last_name}\nEmail: {email}\nRole: {role}\nSalary: {salary}"
        )
    }
}

impl Embed for Record {
    fn embed(
        &self,
        embedder: &mut rig::embeddings::TextEmbedder,
    ) -> Result<(), rig::embeddings::EmbedError> {
        Ok(embedder.embed(self.to_string()))
    }
}

Embedding our CSV

Thankfully, embedding our CSV is actually quite easy now that we've done the hard part. We can create an OpenAI client and get the embedding model, then iterate through each record and build an embedding from it and store it in Rig's in-memory vector store.

You can see below the code to generate an OpenAI client is fairly simple:

let openai_client = rig::providers::openai::Client::from_env();

let agent = openai_client.agent("gpt-4o")
    .preamble("You are a helpful assistant. 
        Your job is to answer a user's questions based on 
        the context given.").build();

// Create embedding model
let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

We'll also keep the code for loading the CSV in as a separate function as it's relatively long.

async fn load_csv(
    model: rig::providers::openai::EmbeddingModel,
) -> Result<
    InMemoryVectorIndex<rig::providers::openai::EmbeddingModel, Record>,
    Box<dyn std::error::Error>,
> {
    let reader = Reader::from_path("employees.csv")?;

    // this iterator attempts to deserialize all rows into the Record type
    // any failed rows will be skipped over rather than causing a panic
    // due to use of filter_map and x.ok() 
    let documents: Vec<Record> = reader
        .into_deserialize::<Record>()
        .filter_map(|x| x.ok()) 
        .collect();

    let documents = match EmbeddingsBuilder::new(model.clone())
        .documents(documents)?
        .build()
        .await
    {
        Ok(ok) => ok,
        Err(e) => return Err(format!("Got error while embedding: {e}").into()),
    };

    let vector_store = InMemoryVectorStore::from_documents(documents);

    Ok(vector_store.index(model))
}

Talking to our CSV file

The next step will be to actually take input from the user and prompt the LLM using not only the prompt, but additional (relevant!) snippets from the vector store.

We'll start by taking some input from standard input (i.e., you load up the application in the terminal, then you type in something and press Enter). We then check if the prompt is equal to quit - if it is then immediately break the loop and close the application, if not, carry on and retrieve a response from the LLM.

println!(
        "Hi! This is your PDF ragger. Write a prompt and press Enter or write \"quit\" to exit."
    );

loop {
    print!("> ");
    let stdin = std::io::stdin();
    let mut prompt = String::new();
    io::stdout().flush()?;
    stdin.read_line(&mut prompt)?;

    match prompt.trim() {
        "quit" => break,
        _ => {}
    }

    let docs = index.top_n::<Record>(prompt.trim(), 4).await?;

    let prompt = format!(
        "Relevant employees:\n{}\n\n{}",
        docs.into_iter()
            .map(|(_, _, doc)| doc.to_string())
            .collect::<Vec<String>>()
            .join("\n"),
        prompt,
    );

    let answer = agent.prompt(prompt.as_ref()).await?;

    println!("{answer}");
}

Creating a chat loop

In terms of storing your conversation history, it is actually quite simple to do so. All you need to do is to create a Vec<Message> and then feed it into the agent, using the chat function instead of prompt. You then add a user message and assistant message (in that order) at the end of the loop iteration. This allows you to build up a conversation with the LLM.

Additionally, we'll also add a reset command which allows the entire conversation to be wiped should the user want to ask the LLM about something else.

println!(
        "Hi! This is your PDF ragger. Write a prompt and press Enter or write \"quit\" to exit. Alternatively, use \"reset\" to reset the conversation."
    );

let mut chat_history: Vec<Message> = Vec::new();

loop {
    print!("> ");
    let stdin = std::io::stdin();
    let mut prompt = String::new();
    io::stdout().flush()?;
    stdin.read_line(&mut prompt)?;

    match prompt.trim() {
        "quit" => break,
        "reset" => {
            chat_history.clear();
            println!("Your conversation has been reset.");
            continue;
        }
        _ => {}
    }

    let docs = index.top_n::<Record>(prompt.trim(), 4).await?;

    let prompt = format!(
        "Relevant employees:\n{}\n\n{}",
        docs.into_iter()
            .map(|(_, _, doc)| doc.to_string())
            .collect::<Vec<String>>()
            .join("\n"),
        prompt,
    );

    let answer = agent.chat(prompt.as_ref(), chat_history.clone()).await?;

    println!("{answer}");
    chat_history.push(Message::user(prompt));
    chat_history.push(Message::assistant(answer));
}

Beyond this tutorial

Now that we've set up a base for your command-line tool, here are a few ways you can extend this demo to do much more than just simple ragging:

• Try adding some tools to your agent!
• What about semantic routing to make sure your LLM stays on track?
• Try incorporating CSV RAG into a new, bigger pipeline!

Conclusion

Thanks for reading! Hopefully you will have a good example of how to improve how you work with data by creating a RAG pipeline from your CSV.