MongoDB Atlas + Fireworks AI RAG Example

python

!pip install -q llama-index llama-index-vector-stores-mongodb llama-index-embeddings-fireworks==0.1.2 llama-index-llms-fireworks
!pip install -q pymongo datasets pandas

python

# set up Fireworks.ai Key
import os
import getpass

fw_api_key = getpass.getpass("Fireworks API Key:")
os.environ["FIREWORKS_API_KEY"] = fw_api_key

python

from datasets import load_dataset
import pandas as pd

# https://huggingface.co/datasets/AIatMongoDB/whatscooking.restaurants
dataset = load_dataset("AIatMongoDB/whatscooking.restaurants")

# Convert the dataset to a pandas dataframe
dataset_df = pd.DataFrame(dataset["train"])

dataset_df.head(5)

python

from llama_index.core.settings import Settings
from llama_index.llms.fireworks import Fireworks
from llama_index.embeddings.fireworks import FireworksEmbedding

embed_model = FireworksEmbedding(
    embed_batch_size=512,
    model_name="nomic-ai/nomic-embed-text-v1.5",
    api_key=fw_api_key,
)
llm = Fireworks(
    temperature=0,
    model="accounts/fireworks/models/mixtral-8x7b-instruct",
    api_key=fw_api_key,
)

Settings.llm = llm
Settings.embed_model = embed_model

python

import json
from llama_index.core import Document
from llama_index.core.schema import MetadataMode

# Convert the DataFrame to a JSON string representation
documents_json = dataset_df.to_json(orient="records")
# Load the JSON string into a Python list of dictionaries
documents_list = json.loads(documents_json)

llama_documents = []

for document in documents_list:
    # Value for metadata must be one of (str, int, float, None)
    document["name"] = json.dumps(document["name"])
    document["cuisine"] = json.dumps(document["cuisine"])
    document["attributes"] = json.dumps(document["attributes"])
    document["menu"] = json.dumps(document["menu"])
    document["borough"] = json.dumps(document["borough"])
    document["address"] = json.dumps(document["address"])
    document["PriceRange"] = json.dumps(document["PriceRange"])
    document["HappyHour"] = json.dumps(document["HappyHour"])
    document["review_count"] = json.dumps(document["review_count"])
    document["TakeOut"] = json.dumps(document["TakeOut"])
    # these two fields are not relevant to the question we want to answer,
    # so I will skip it for now
    del document["embedding"]
    del document["location"]

    # Create a Document object with the text and excluded metadata for llm and embedding models
    llama_document = Document(
        text=json.dumps(document),
        metadata=document,
        metadata_template="{key}=>{value}",
        text_template="Metadata: {metadata_str}\n-----\nContent: {content}",
    )

    llama_documents.append(llama_document)

# Observing an example of what the LLM and Embedding model receive as input
print(
    "\nThe LLM sees this: \n",
    llama_documents[0].get_content(metadata_mode=MetadataMode.LLM),
)
print(
    "\nThe Embedding model sees this: \n",
    llama_documents[0].get_content(metadata_mode=MetadataMode.EMBED),
)

python

llama_documents[0]

python

from llama_index.core.node_parser import SentenceSplitter

parser = SentenceSplitter()
nodes = parser.get_nodes_from_documents(llama_documents)
# 25k nodes takes about 10 minutes, will trim it down to 2.5k
new_nodes = nodes[:2500]

# There are 25k documents, so we need to do batching. Fortunately LlamaIndex provides good batching
# for embedding models, and we are going to rely on the __call__ method for the model to handle this
node_embeddings = embed_model(new_nodes)

python

for idx, n in enumerate(new_nodes):
    n.embedding = node_embeddings[idx].embedding
    if "_id" in n.metadata:
        del n.metadata["_id"]

Ensure your database, collection and vector store index is setup on MongoDB Atlas for the collection or the following step won't work appropriately on MongoDB.

For assistance with database cluster setup and obtaining the URI, refer to this guide for setting up a MongoDB cluster, and this guide to get your connection string.
Once you have successfully created a cluster, create the database and collection within the MongoDB Atlas cluster by clicking “+ Create Database”. The database will be named movies, and the collection will be named movies_records.
Creating a vector search index within the movies_records collection is essential for efficient document retrieval from MongoDB into our development environment. To achieve this, refer to the official guide on vector search index creation.

python

import pymongo


# set up Fireworks.ai Key
import os
import getpass

mongo_uri = getpass.getpass("MONGO_URI:")
if not mongo_uri:
    print("MONGO_URI not set")

mongo_client = pymongo.MongoClient(mongo_uri)
async_mongo_client = pymongo.AsyncMongoClient(mongo_uri)

DB_NAME = "whatscooking"
COLLECTION_NAME = "restaurants"

db = mongo_client[DB_NAME]
collection = db[COLLECTION_NAME]

python

# To ensure we are working with a fresh collection
# delete any existing records in the collection
collection.delete_many({})

python

from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch

vector_store = MongoDBAtlasVectorSearch(
    mongodb_client=mongo_client,
    async_mongodb_client=async_mongo_client,
    db_name=DB_NAME,
    collection_name=COLLECTION_NAME,
    index_name="vector_index",
)
vector_store.add(new_nodes)

now make sure you create the search index with the right name here

python

from llama_index.core import VectorStoreIndex, StorageContext

index = VectorStoreIndex.from_vector_store(vector_store)

python

%pip install -q matplotlib

python

import pprint
from llama_index.core.response.notebook_utils import display_response

query_engine = index.as_query_engine()

query = "search query: Anything that doesn't have alcohol in it"

response = query_engine.query(query)
display_response(response)
pprint.pprint(response.source_nodes)