Code Summary: Retrieving the Data

This code implements the retriever component of our RAG system, using Atlas Vector Search. You can also view or fork the code from the Curriculum GitHub repository.

Prerequisites

• Atlas Cluster Connection String
• OpenAI API Key
• Data loaded into Atlas

Create the following vector search index, named vector_index, on the chunked_data collection in your Atlas Cluster:

{
  "fields": [
    {
      "numDimensions": 1536,
      "path": "embedding",
      "similarity": "cosine",
      "type": "vector"
    },
    {
      "path": "hasCode",
      "type": "filter"
    }
  ]
}

Usage

rag.py file

The following code sets Atlas Vector Search as the vectorStore. It takes the query embeddings as input and uses that to perform a vector search to retrieve relevant data.

from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings
import key_param

dbName = "book_mongodb_chunks"
collectionName = "chunked_data"
index = "vector_index"

vectorStore = MongoDBAtlasVectorSearch.from_connection_string(
    key_param.MONGODB_URI,
    dbName + "." + collectionName,
    OpenAIEmbeddings(disallowed_special=(), openai_api_key=key_param.LLM_API_KEY),
    index_name=index,
)

def query_data(query):
    retriever = vectorStore.as_retriever(
        search_type="similarity",
        search_kwargs={
            "k": 3
        },
    )

    results = retriever.invoke(query)
    print(results)

Update the query_data function in the rag.py file by passing in a question:

query_data("What is the difference between a database and collection in MongoDB?")

Add a prefilter to the retriever

Update the query_data function in the rag.py file to include a pre_filter object.

def query_data(query):
    retriever = vectorStore.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={
            "k": 3,
            "pre_filter": { "hasCode": { "$eq": False } },
            "score_threshold": 0.01
        },
    )
    results = retriever.invoke(query)
    print(results)

rag.py file with filtering

This code generates vector embeddings for a query and uses them to perform a vector search on the chunked_data collection. It also prefilters the data to find chunks without code.

from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings
import key_param

dbName = "book_mongodb_chunks"
collectionName = "chunked_data"
index = "vector_index"

vectorStore = MongoDBAtlasVectorSearch.from_connection_string(
    key_param.MONGODB_URI,
    dbName + "." + collectionName,
    OpenAIEmbeddings(disallowed_special=(), openai_api_key=key_param.LLM_API_KEY),
    index_name=index,
)

def query_data(query):
    retriever = vectorStore.as_retriever(
        search_type="similarity",
        search_kwargs={
            "k": 3,
            "pre_filter": { "hasCode": { "$eq": False } },
            "score_threshold": 0.01
        },
    )

    results = retriever.invoke(query)
    print(results)

    

query_data("When did MongoDB begin supporting multi-document transactions?")

Run the demo:

python rag.py