Blog

Engineering

End-to-end RAG with Airbyte Cloud, Google Drive, and PGVector

Build an end-to-end RAG pipeline with Airbyte Cloud, Google Drive, and pgvector, syncing documents into Postgres for semantic search.

ENGINEERING

October 21, 2024

10 min read

Aldo Gonzalez

Summarize with AI:

End-to-end RAG with Airbyte Cloud, Google Drive, and PGVector

Have you ever wondered how it could remember specific episode numbers, seasons, and quotes or your favorite shows like those comic store experts, "Oh yes, season 9, episode 10, when Homer saves Mark Hamil…"? We can’t pull all that data in our brain right now, but what if you could build an AI assistant for this by using Airbyte?

‍

With the rapid pace of AI tools, building an assistant is easier than you think. In this post, we'll walk through using Airbyte to pass a document with quotes and episode information from a top-rated TV show (with yellow people) to PGVector for processing. This process, called retrieval-augmented generation or RAG, leverages PGvector functions to consume and analyze the data seamlessly. Good news is we just released a PGVector connector to support RAG use cases just like this.

TL;DR:

We will move Simpsons data, sourced from Kaggle, and load it into PGvector, which allows us to perform a cosine similarity search using OpenAI.

‍

Understanding RAG

LLMs, or Large Language Models, are quite helpful when you are interested in general information, and frequently used as the source of data for AI. For tasks requiring deeper insights or specific analysis, relying on the AI research assistant can significantly streamline the processUnfortunately, LLMs far short regarding domain-specific information, and that's where they start hallucinating. A hallucination is basically a wrong or inaccurate response to your prompt. To address this limitation the LLM must be provided with up-to-date information from an accurate data source. This process is called Retrieval-Augmented Generation (RAG). Basically, you can’t have AI without data. And without RAG, the data can’t always be trusted.

Prerequisites

1. Data Source: In this tutorial we use a Google Drive folder. I suggest you download these files and load into your own Google drive to ensure you have a copy of it. Note: You will find a "denormalized" csv named episode_spoken_words.csv with reduced data for this guide's purposes. You can use the other files from which I created that one. All the other CSVs were downloaded from Kaggle.

2. Airbyte Cloud Account : Log in or sign up here. You can also use OSS.

3. Postgres DB with PGVector extension installed. If you need guidance, check out the PGVector documentation online.

4. OpenAI: You will need an OpenAI account and api key for authentication. Also, Ensure you do not have a rate limit to continue.

Steps

1. Setup Airbyte Data Source

In Airbyte, tap on Sources , then select the Google Drive connector as your source. Paste the folder URL to episode_spoken_words.csv, and create a stream with the CSV Format (Experimental). Finally, test it to ensure it's perfectly set up.

‍

2. Setup PGVector Destination

To set up a PGVector instance, you need to set up Postgres entities (database, schema, user, role, and extensions), as explained in this documentation.

You can run the following SQL commands, but there are different ways to achieve this. Don’t forget to change your_password_here with your actual Airbyte password.


-- Create a database if you don't have one
create database pgvector_simpsons_tutorial;

-- The following commands will create a new user:
CREATE USER airbyte_user PASSWORD 'your_password_here';

-- Now, provide this user with read-only access to relevant schemas and tables. Re-run this command for each schema you expect to replicate data from:

GRANT USAGE ON SCHEMA your_schema_name TO airbyte_user;
GRANT SELECT ON ALL TABLES IN SCHEMA your_schema_name TO airbyte_user;
ALTER DEFAULT PRIVILEGES IN SCHEMA your_schema_name 
GRANT SELECT ON TABLES TO airbyte_user;

-- Enable the extension 
CREATE EXTENSION vector;

Once run, this script will spin up a database, ready to store your data. Move back to Airbyte, and tap Destination, and select PGVector. Configure using the following values. Save by tapping Set up destination. Finally, test the source to make sure it's working as expected.

Processing Chunk size: Different embedding models have different token limitations. In this tutorial I used 1000 for Open AI embedding option.The best chunking is also dependent on the data you are dealing with.
Embedding: Paste your OpenAI api key and save.
Indexes: Set up as below

‍

3. Move Data

Next, we need to create a connection between your Google Drive source and the PGVector destination you just created. This will sync the data into PGVector. . Here is an example of successful connection after sync.

‍

4: Explore data in Postgres

At this point you should be able to see the data in Postgres. The following columns will be available in your database, tables have the following columns

DOCUMENT_ID - unique based on primary key
CHUNK_ID - randomly generated uuid
DOCUMENT_CONTENT - text context from source
METADATA - metadata from source
EMBEDDING - Vector representation of the document_content

Here is a snippet of how one of my results appears.

‍

5: Building the RAG with PGVector Functions.

Congratulations for making it this far. You’ve done the hard work. Now it is time to see the payoff. First, a little RAG theory to help you. RAG heavily relies on semantic comparison techniques. The measurement of similarity between vectors is a fundamental operation in semantic comparison. This operation is used to find the top N closest vectors to a query vector, which can be used for semantic search. Basically, vector search enables developers to improve the accuracy of their generative AI responses by providing related documents to a large language model. Ok, back to the tutorial. Lesson over.

‍

The key elements in the RAG process we have to implement are :

Generate Embeddings from the query : Converting a question into a vector array.

Ideally, you can either embed data using OpenAI, Cohere, an OpenAI compatible source, or Faker. Once you have the embedded data, you need to embed the question with the appropriate method distinct from each model.

In our example, we are using OpenAI to help us out. To generate the embeddings using OpenAI you can use the embeddings functions.

‍

Breaking it down a bit, we are performing what is called a similarity search to find matching chunks of data. The relevance, or how accurate a result is to the question asked is referred to as distance. In other words, the closer the distance, the more accurate the answer. Once we have the question vector or embedding, we can then create an SQL statement to query the database. Within SQL there are a number of different ways to measure the distance. For reference, I’ve listed them below and encourage you to do more research on your own for future projects. For our need, L2 distance works great.

PGVector supported distance functions are:

<-> - L2 distance: We will use this function in our demo

<#> - (negative) inner product

<=> - cosine distance

<+> - L1 distance (added in 0.7.0)

<~> - Hamming distance (binary vectors, added in 0.7.0)

<%> - Jaccard distance (binary vectors, added in 0.7.0)

With our question vector and understanding of distance functions, we can write a query and function to get the related context from Postgres:


engine = create_engine(url)
QUERY_TEMPLATE = """
SELECT document_content, 
    metadata->'title' as episode_title, 
    metadata->'script_line_number' as script_line_number,
    metadata->'name' as character_name,
    metadata->'spoken_words' as spoken_words
FROM episode_spoken_words
ORDER BY embedding <-> '{question_vector}'
LIMIT 5
"""

def get_context(question) -> str:
    question_vector = get_embedding_from_open_ai(question)
    query = QUERY_TEMPLATE.format(question_vector=question_vector)
    result = engine.execute(query)
    return ("\n\n" + "-" * 8 + "\n\n").join([f"Episode {row.episode_title} | Line number: {row.script_line_number} | Spoken Words: {row.spoken_words} | Character: {row.character_name}" for row in result])

Finally, we get the response, and use Open AI LLM for our agent:

‍

With everything in place, you can ask the RAG questions about the Simpsons, free of hallucinations. Pretty cool, huh?

‍

Summary

This post provided a step-by-step guide on leveraging Airbyte to sync unstructured data into PGVector for RAG. We also used OpenAI as the LLM. Throughout, we tried to provide some broader explanations of concepts important to help you better understand how to create your own AI agents. As we saw in our demo, the measurement of similarity between vectors is a fundamental operation in semantic comparison and providing hallucination free results, even when working with critical datasets such as the Simpsons. If you are interested in building a similar solution you can sign up for Airbyte Cloud or install Airbyte using abctl to get started. As Bart Simspon would say, "¡Ay, caramba!", that’s it.

About the Author

Aldo Gonzalez