LlamaIndexTS retrieval augmented generation with MongoDB

Prepare Environment

Read and follow the instructions in the README.md file located one directory up to make sure your JS/TS dependencies are set up. The commands listed below are also run from that parent directory.

Sign up for MongoDB Atlas

We'll be using MongoDB's hosted database service, MongoDB Atlas. You can sign up for free and get a small hosted cluster for free:

The signup process will walk you through the process of creating your cluster and ensuring it's configured for you to access. Once the cluster is created, choose "Connect" and then "Connect to your application". Choose Python, and you'll be presented with a connection string that looks like this:

Set up environment variables

Copy the connection string (make sure you include your password) and put it into a file called .env in the parent folder of this directory. It should look like this:

MONGODB_URI=mongodb+srv://seldo:xxxxxxxxxxx@llamaindexdemocluster.xfrdhpz.mongodb.net/?retryWrites=true&w=majority

You will also need to choose a name for your database, and the collection where we will store the tweets, and also include them in .env. They can be any string, but this is what we used:

MONGODB_DATABASE=tiny_tweets_db
MONGODB_COLLECTION=tiny_tweets_collection

Import tweets into MongoDB

You are now ready to import our ready-made data set into Mongo. This is the file tinytweets.json, a selection of approximately 1000 tweets from @seldo on Twitter in mid-2019. With your environment set up you can do this by running

npx tsx mongodb/1_import.ts

If you don't want to use tweets, you can replace json_file with any other array of JSON objects, but you will need to modify some code later to make sure the correct field gets indexed. There is no LlamaIndex-specific code here; you can load your data into Mongo any way you want to.

Load and index your data

Now we're ready to index our data. To do this, LlamaIndex will pull your text out of Mongo, split it into chunks, and then send those chunks to OpenAI to be turned into vector embeddings. The embeddings will then be stored in a new collection in Mongo. This will take a while depending how much text you have, but the good news is that once it's done you will be able to query quickly without needing to re-index.

We'll be using OpenAI to do the embedding, so now is when you need to generate an OpenAI API key if you haven't already and add it to your .env file like this:

OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

You'll also need to pick a name for the new collection where the embeddings will be stored, and add it to .env, along with the name of a vector search index (we'll be creating this in the next step, after you've indexed your data):

MONGODB_VECTORS=tiny_tweets_vectors
MONGODB_VECTOR_INDEX=tiny_tweets_vector_index

If the data you're indexing is the tweets we gave you, you're ready to go:

npx tsx mongodb/2_load_and_index.ts

Note: this script is running a couple of minutes and currently doesn't show any progress.

What you're doing here is creating a Reader which loads the data out of Mongo in the collection and database specified. It looks for text in a set of specific keys in each object. In this case we've given it just one key, "full_text".

Now you're creating a vector search client for Mongo. In addition to a MongoDB client object, you again tell it what database everything is in. This time you give it the name of the collection where you'll store the vector embeddings, and the name of the vector search index you'll create in the next step.

Run a test query

You can do this by running

npx tsx mongodb/3_query.ts

This sets up a connection to Atlas just like 2_load_and_index.ts did, then it creates a query engine and runs a query against it.

If all is well, you should get a nuanced opinion about web frameworks.