Optimizing for Relevance Using MongoDB Atlas and LlamaIndex

1
!pip install -qU pymongo llama-index llama-index-llms-openai llama-index-vector-stores-mongodb

1
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter your OpenAI API key: ")

1
MONGODB_URI = getpass.getpass("Enter your MongoDB URI: ")
2
mongodb_client = MongoClient(
3
    MONGODB_URI, appname="devrel.content.retrieval_strategies_llamaindex"
4
)

1
# Load the dataset
2
data = load_dataset("MongoDB/embedded_movies", split="train")
3
data = pd.DataFrame(data)
4
5
# Fill Nones in the dataframe
6
data = data.fillna({"genres": "[]", "languages": "[]", "cast": "[]", "imdb": "{}"})

1
documents = []
2
for _, row in data.iterrows():
3
    # Extract required fields
4
    title = row["title"]
5
    rating = row["imdb"].get("rating", 0)
6
    languages = row["languages"]
7
    cast = row["cast"]
8
    genres = row["genres"]
9
    # Create the metadata attribute
10
    metadata = {"title": title, "rating": rating, "languages": languages}
11
    # Create the text attribute
12
    text = f"Title: {title}\nPlot: {row['fullplot']}\nCast: {', '.join(item for item in cast)}\nGenres: {', '.join(item for item in  genres)}\nLanguages: {', '.join(item for item in languages)}\nRating: {rating}"
13
    documents.append(Document(text=text, metadata=metadata))

1
Document(id_='df580995-076f-444e-ae21-9ba596a6a224', embedding=None, metadata={'title': 'The Perils of Pauline', 'rating': 7.6, 'languages': ['English']}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, text='Title: The Perils of Pauline\nPlot: Young Pauline is left a lot of money...\nCast: Pearl White, Crane Wilbur, Paul Panzer, Edward Josè\nGenres: Action\nLanguages: English\nRating: 7.6', mimetype='text/plain', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\n\n{content}', metadata_template='{key}: {value}', metadata_seperator='\n')

1
from llama_index.embeddings.openai import OpenAIEmbedding
2
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
3
from llama_index.core.settings import Settings
4
from llama_index.core import VectorStoreIndex, StorageContext
5
6
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
7
8
VS_INDEX_NAME = "vector_index"
9
FTS_INDEX_NAME = "fts_index"
10
DB_NAME = "llamaindex"
11
COLLECTION_NAME = "hybrid_search"
12
collection = mongodb_client[DB_NAME][COLLECTION_NAME]
13
14
vector_store = MongoDBAtlasVectorSearch(
15
    mongodb_client,
16
    db_name=DB_NAME,
17
    collection_name=COLLECTION_NAME,
18
    vector_index_name=VS_INDEX_NAME,
19
    fulltext_index_name=FTS_INDEX_NAME,
20
    embedding_key="embedding",
21
    text_key="text",
22
)
23
# If the collection has documents with embeddings already, create the vector store index from the vector store
24
if collection.count_documents({}) > 0:
25
    vector_store_index = VectorStoreIndex.from_vector_store(vector_store)
26
# If the collection does not have documents, embed and ingest them into the vector store
27
else:
28
    vector_store_context = StorageContext.from_defaults(vector_store=vector_store)
29
    vector_store_index = VectorStoreIndex.from_documents(
30
        documents, storage_context=vector_store_context, show_progress=True
31
    )

1
vs_model = SearchIndexModel(
2
    definition={
3
        "fields": [
4
            {
5
                "type": "vector",
6
                "path": "embedding",
7
                "numDimensions": 1536,
8
                "similarity": "cosine",
9
            },
10
            {"type": "filter", "path": "metadata.rating"},
11
            {"type": "filter", "path": "metadata.language"},
12
        ]
13
    },
14
    name=VS_INDEX_NAME,
15
    type="vectorSearch",
16
)

1
fts_model = SearchIndexModel(
2
    definition={"mappings": {"dynamic": False, "fields": {"text": {"type": "string"}}}},
3
    name=FTS_INDEX_NAME,
4
    type="search",
5
)

1
for model in [vs_model, fts_model]:
2
    try:
3
        collection.create_search_index(model=model)
4
    except OperationFailure:
5
        print(f"Duplicate index found for model {model}. Skipping index creation.")

1
def get_recommendations(query: str, mode: str, **kwargs) -> None:
2
    """
3
    Get movie recommendations
4
5
    Args:
6
        query (str): User query
7
        mode (str): Retrieval mode. One of (default, text_search, hybrid)
8
    """
9
    query_engine = vector_store_index.as_query_engine(
10
        similarity_top_k=5, vector_store_query_mode=mode, **kwargs
11
    )
12
    response = query_engine.query(query)
13
    nodes = response.source_nodes
14
    for node in nodes:
15
        title = node.metadata["title"]
16
        rating = node.metadata["rating"]
17
        score = node.score
18
        print(f"Title: {title} | Rating: {rating} | Relevance Score: {score}")

1
get_recommendations(
2
    query="Action movies about humans fighting machines",
3
    mode="text_search",
4
)
5
6
Title: Hellboy II: The Golden Army | Rating: 7.0 | Relevance Score: 5.93734884262085
7
Title: The Matrix Revolutions | Rating: 6.7 | Relevance Score: 4.574477195739746
8
Title: The Matrix | Rating: 8.7 | Relevance Score: 4.387373924255371
9
Title: Go with Peace Jamil | Rating: 6.9 | Relevance Score: 3.5394840240478516
10
Title: Terminator Salvation | Rating: 6.7 | Relevance Score: 3.3378987312316895

1
get_recommendations(
2
    query="Action movies about humans fighting machines", mode="default"
3
)
4
5
Title: Death Machine | Rating: 5.7 | Relevance Score: 0.7407287359237671
6
Title: Real Steel | Rating: 7.1 | Relevance Score: 0.7364246845245361
7
Title: Soldier | Rating: 5.9 | Relevance Score: 0.7282171249389648
8
Title: Terminator 3: Rise of the Machines | Rating: 6.4 | Relevance Score: 0.7266112565994263
9
Title: Last Action Hero | Rating: 6.2 | Relevance Score: 0.7250100374221802

1
get_recommendations(query="Action movies about humans fighting machines", mode="hybrid")
2
3
Title: Hellboy II: The Golden Army | Rating: 7.0 | Relevance Score: 0.5
4
Title: Death Machine | Rating: 5.7 | Relevance Score: 0.5
5
Title: The Matrix Revolutions | Rating: 6.7 | Relevance Score: 0.25
6
Title: Real Steel | Rating: 7.1 | Relevance Score: 0.25
7
Title: Soldier | Rating: 5.9 | Relevance Score: 0.16666666666666666

1
# Higher alpha, vector search dominates
2
get_recommendations(
3
    query="Action movies about humans fighting machines",
4
    mode="hybrid",
5
    alpha=0.7,
6
)
7
8
Title: Death Machine | Rating: 5.7 | Relevance Score: 0.7
9
Title: Real Steel | Rating: 7.1 | Relevance Score: 0.35
10
Title: Hellboy II: The Golden Army | Rating: 7.0 | Relevance Score: 0.30000000000000004
11
Title: Soldier | Rating: 5.9 | Relevance Score: 0.2333333333333333
12
Title: Terminator 3: Rise of the Machines | Rating: 6.4 | Relevance Score: 0.175

1
# Lower alpha, full-text search dominates
2
get_recommendations(
3
    query="Action movies about humans fighting machines",
4
    mode="hybrid",
5
    alpha=0.3,
6
)
7
8
Title: Hellboy II: The Golden Army | Rating: 7.0 | Relevance Score: 0.7
9
Title: The Matrix Revolutions | Rating: 6.7 | Relevance Score: 0.35
10
Title: Death Machine | Rating: 5.7 | Relevance Score: 0.3
11
Title: The Matrix | Rating: 8.7 | Relevance Score: 0.2333333333333333
12
Title: Go with Peace Jamil | Rating: 6.9 | Relevance Score: 0.175

1
from llama_index.core.vector_stores import (
2
    MetadataFilter,
3
    MetadataFilters,
4
    FilterOperator,
5
    FilterCondition,
6
)
7
8
filters = MetadataFilters(
9
    filters=[
10
        MetadataFilter(key="metadata.rating", value=7, operator=FilterOperator.GT),
11
        MetadataFilter(
12
            key="metadata.languages", value="English", operator=FilterOperator.EQ
13
        ),
14
    ],
15
    condition=FilterCondition.AND,
16
)

1
get_recommendations(
2
    query="Action movies about humans fighting machines",
3
    mode="hybrid",
4
    alpha=0.7,
5
    filters=filters,
6
)
7
8
Title: Real Steel | Rating: 7.1 | Relevance Score: 0.7
9
Title: T2 3-D: Battle Across Time | Rating: 7.8 | Relevance Score: 0.35
10
Title: The Matrix | Rating: 8.7 | Relevance Score: 0.30000000000000004
11
Title: Predator | Rating: 7.8 | Relevance Score: 0.2333333333333333
12
Title: Transformers | Rating: 7.1 | Relevance Score: 0.175