trying out MongoDB vector search using OpenAI

setup Python environment in Jupyter notebook,

%pip install --upgrade --quiet  langchain langchain_community langchain_core langchain_openai pymongo lark

import os

OPENAI_API_KEY = "sk-xxxxx"
os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

add current IP address to MongoDB,

!curl ipecho.net/plain

then load the data to MongoDB,

from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from pymongo import MongoClient

CONNECTION_STRING = "mongodb+srv://rw-user-1:$pwd@clusterm0.nnnnn.mongodb.net/?tls=true&retryWrites=true&w=majority&appName=ClusterM0"
DB_NAME = "media"
COLLECTION_NAME = "movie"
INDEX_NAME = "default"

MongoClient = MongoClient(CONNECTION_STRING)
collection = MongoClient[DB_NAME][COLLECTION_NAME]

embeddings = OpenAIEmbeddings()

docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={"year": 1993, "rating": 7.7, "genre": "action"},
    ),
    Document(
        page_content="Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
        metadata={"year": 2010, "genre": "thriller", "rating": 8.2},
    ),
    Document(
        page_content="A bunch of normal-sized women are supremely wholesome and some men pine after them",
        metadata={"year": 2019, "rating": 8.3, "genre": "drama"},
    ),
    Document(
        page_content="Three men walk into the Zone, three men walk out of the Zone",
        metadata={"year": 1979, "rating": 9.9, "genre": "science fiction"},
    ),
    Document(
        page_content="A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea",
        metadata={"year": 2006, "genre": "thriller", "rating": 9.0},
    ),
    Document(
        page_content="Toys come alive and have a blast doing so",
        metadata={"year": 1995, "genre": "animated", "rating": 9.3},
    ),
]

vectorstore = MongoDBAtlasVectorSearch.from_documents(
    docs,
    embeddings,
    collection=collection,
    index_name=INDEX_NAME,
)

now, we need to add in the search index to the database,

{
  "mappings": {
    "dynamic": true,
    "fields": {
      "embedding": {
        "dimensions": 1536,
        "similarity": "cosine",
        "type": "knnVector"
      },
      "genre": {
        "type": "token"
      },
      "ratings": {
        "type": "number"
      },
      "year": {
        "type": "number"
      }
    }
  }
}

now let’s run a search

def create_vector_search():
    """
    Creates a MongoDBAtlasVectorSearch object using the connection string, database, and collection names, along with the OpenAI embeddings and index configuration.

    :return: MongoDBAtlasVectorSearch object
    """
    vector_search = MongoDBAtlasVectorSearch.from_connection_string(
        CONNECTION_STRING,
        f"{DB_NAME}.{COLLECTION_NAME}",
        embeddings,
        index_name=INDEX_NAME
    )
    return vector_search

def perform_similarity_search(query, top_k=1):
    """
    This function performs a similarity search within a MongoDB Atlas collection. It leverages the capabilities of the MongoDB Atlas Search, which under the hood, may use the `$vectorSearch` operator, to find and return the top `k` documents that match the provided query semantically.

    :param query: The search query string.
    :param top_k: Number of top matches to return.
    :return: A list of the top `k` matching documents with their similarity scores.
    """

   # Get the MongoDBAtlasVectorSearch object
    vector_search = create_vector_search()

    # Execute the similarity search with the given query
    results = vector_search.similarity_search_with_score(
        query=query,
        k=top_k,
    )

    return results

# Example of calling the function directly
#relevant_documents = perform_similarity_search("What are some movies about dinosaurs")
relevant_documents = perform_similarity_search("I want to watch a movie about toys rated higher than 9")
print("Number of relevant documents:", len(relevant_documents))

for document in relevant_documents:
    print("Document content:", document)
    print()

we get some similar records as following

Number of relevant documents: 1