Atlas quickstart

MongoDB Atlas Quickstart¶

MongoDB Atlas Vector Search is part of the MongoDB platform that enables MongoDB customers to build intelligent applications powered by semantic search over any type of data. Atlas Vector Search allows you to integrate your operational database and vector search in a single, unified, fully managed platform with full vector database capabilities.

You can integrate TruLens with your application built on Atlas Vector Search to leverage observability and measure improvements in your application's search capabilities.

This tutorial will walk you through the process of setting up TruLens with MongoDB Atlas Vector Search and Llama-Index as the orchestrator.

Even better, you'll learn how to use metadata filters to create specialized query engines and leverage a router to choose the most appropriate query engine based on the query.

See MongoDB Atlas/LlamaIndex Quickstart for more details.

In [ ]:

# !pip install trulens trulens-apps-llamaindex trulens-providers-openai llama-index llama-index-vector-stores-mongodb llama-index-embeddings-openai pymongo

# !pip install trulens trulens-apps-llamaindex trulens-providers-openai llama-index llama-index-vector-stores-mongodb llama-index-embeddings-openai pymongo

Import TruLens and start the dashboard¶

In [ ]:

from trulens.core import TruSession
from trulens.dashboard import run_dashboard

session = TruSession()
session.reset_database()
run_dashboard(session)

from trulens.core import TruSession from trulens.dashboard import run_dashboard session = TruSession() session.reset_database() run_dashboard(session)

Set imports, keys and llama-index settings¶

In [ ]:

import os

from llama_index.core import SimpleDirectoryReader
from llama_index.core import StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.settings import Settings
from llama_index.core.vector_stores import ExactMatchFilter
from llama_index.core.vector_stores import MetadataFilters
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
import pymongo

import os from llama_index.core import SimpleDirectoryReader from llama_index.core import StorageContext from llama_index.core import VectorStoreIndex from llama_index.core.query_engine import RetrieverQueryEngine from llama_index.core.retrievers import VectorIndexRetriever from llama_index.core.settings import Settings from llama_index.core.vector_stores import ExactMatchFilter from llama_index.core.vector_stores import MetadataFilters from llama_index.embeddings.openai import OpenAIEmbedding from llama_index.llms.openai import OpenAI from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch import pymongo

In [ ]:

os.environ["OPENAI_API_KEY"] = "sk-..."
ATLAS_CONNECTION_STRING = (
    "mongodb+srv://<username>:<password>@<clusterName>.<hostname>.mongodb.net"
)

os.environ["OPENAI_API_KEY"] = "sk-..." ATLAS_CONNECTION_STRING = ( "mongodb+srv://:@..mongodb.net" )

In [ ]:

Settings.llm = OpenAI()
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")
Settings.chunk_size = 100
Settings.chunk_overlap = 10

Settings.llm = OpenAI() Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002") Settings.chunk_size = 100 Settings.chunk_overlap = 10

Load sample data¶

Here we'll load two PDFs: one for Atlas best practices and one textbook on database essentials.

In [ ]:

# Load the sample data
!mkdir -p 'data/'
!wget 'https://query.prod.cms.rt.microsoft.com/cms/api/am/binary/RE4HkJP' -O 'data/atlas_best_practices.pdf'
atlas_best_practices = SimpleDirectoryReader(
    input_files=["./data/atlas_best_practices.pdf"]
).load_data()

!wget 'http://fondamentidibasididati.it/wp-content/uploads/2020/11/DBEssential-2021-C30-11-21.pdf' -O 'data/DBEssential-2021.pdf'
db_essentials = SimpleDirectoryReader(
    input_files=["./data/DBEssential-2021.pdf"]
).load_data()

!wget 'https://courses.edx.org/asset-v1:Databricks+LLM101x+2T2023+type@asset+block@Module_2_slides.pdf' -O 'data/DataBrick_vector_search.pdf'
databrick_vector_search = SimpleDirectoryReader(
    input_files=["./data/DataBrick_vector_search.pdf"]
).load_data()

documents = atlas_best_practices + db_essentials + databrick_vector_search

# Load the sample data !mkdir -p 'data/' !wget 'https://query.prod.cms.rt.microsoft.com/cms/api/am/binary/RE4HkJP' -O 'data/atlas_best_practices.pdf' atlas_best_practices = SimpleDirectoryReader( input_files=["./data/atlas_best_practices.pdf"] ).load_data() !wget 'http://fondamentidibasididati.it/wp-content/uploads/2020/11/DBEssential-2021-C30-11-21.pdf' -O 'data/DBEssential-2021.pdf' db_essentials = SimpleDirectoryReader( input_files=["./data/DBEssential-2021.pdf"] ).load_data() !wget 'https://courses.edx.org/asset-v1:Databricks+LLM101x+2T2023+type@asset+block@Module_2_slides.pdf' -O 'data/DataBrick_vector_search.pdf' databrick_vector_search = SimpleDirectoryReader( input_files=["./data/DataBrick_vector_search.pdf"] ).load_data() documents = atlas_best_practices + db_essentials + databrick_vector_search

Create a vector store¶

Next you need to create an Atlas Vector Search Index.

When you do so, use the following in the json editor:

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

In [ ]:

# Connect to your Atlas cluster
mongodb_client = pymongo.MongoClient(ATLAS_CONNECTION_STRING)

# Instantiate the vector store
atlas_vector_search = MongoDBAtlasVectorSearch(
    mongodb_client,
    db_name="atlas-quickstart-demo",
    collection_name="test",
    index_name="vector_index",
)
vector_store_context = StorageContext.from_defaults(
    vector_store=atlas_vector_search
)

# load both documents into the vector store
vector_store_index = VectorStoreIndex.from_documents(
    documents, storage_context=vector_store_context, show_progress=True
)

# Connect to your Atlas cluster mongodb_client = pymongo.MongoClient(ATLAS_CONNECTION_STRING) # Instantiate the vector store atlas_vector_search = MongoDBAtlasVectorSearch( mongodb_client, db_name="atlas-quickstart-demo", collection_name="test", index_name="vector_index", ) vector_store_context = StorageContext.from_defaults( vector_store=atlas_vector_search ) # load both documents into the vector store vector_store_index = VectorStoreIndex.from_documents( documents, storage_context=vector_store_context, show_progress=True )

Setup basic RAG¶

In [ ]:

query_engine = vector_store_index.as_query_engine()

query_engine = vector_store_index.as_query_engine()

Add feedback functions¶

In [ ]:

import numpy as np
from trulens.core import Feedback
from trulens.providers.openai import OpenAI
from trulens.apps.llamaindex import TruLlama

# Initialize provider class
provider = OpenAI()

# select context to be used in feedback. the location of context is app specific.
context = TruLlama.select_context(query_engine)

# Define a groundedness feedback function
f_groundedness = (
    Feedback(
        provider.groundedness_measure_with_cot_reasons, name="Groundedness"
    )
    .on(context.collect())  # collect context chunks into a list
    .on_output()
)

# Question/answer relevance between overall question and answer.
f_answer_relevance = Feedback(
    provider.relevance_with_cot_reasons, name="Answer Relevance"
).on_input_output()
# Context relevance between question and each context chunk.
f_context_relevance = (
    Feedback(
        provider.context_relevance_with_cot_reasons, name="Context Relevance"
    )
    .on_input()
    .on(context)
    .aggregate(np.mean)
)

import numpy as np from trulens.core import Feedback from trulens.providers.openai import OpenAI from trulens.apps.llamaindex import TruLlama # Initialize provider class provider = OpenAI() # select context to be used in feedback. the location of context is app specific. context = TruLlama.select_context(query_engine) # Define a groundedness feedback function f_groundedness = ( Feedback( provider.groundedness_measure_with_cot_reasons, name="Groundedness" ) .on(context.collect()) # collect context chunks into a list .on_output() ) # Question/answer relevance between overall question and answer. f_answer_relevance = Feedback( provider.relevance_with_cot_reasons, name="Answer Relevance" ).on_input_output() # Context relevance between question and each context chunk. f_context_relevance = ( Feedback( provider.context_relevance_with_cot_reasons, name="Context Relevance" ) .on_input() .on(context) .aggregate(np.mean) )

In [ ]:

tru_query_engine_recorder = TruLlama(
    query_engine,
    app_name="RAG",
    app_version="Basic RAG",
    feedbacks=[f_groundedness, f_answer_relevance, f_context_relevance],
)

tru_query_engine_recorder = TruLlama( query_engine, app_name="RAG", app_version="Basic RAG", feedbacks=[f_groundedness, f_answer_relevance, f_context_relevance], )

Write test cases¶

Let's write a few test queries to test the ability of our RAG to answer questions on both documents in the vector store.

In [ ]:

test_set = {
    "MongoDB Atlas": [
        "How do you secure MongoDB Atlas?",
        "How can Time to Live (TTL) be used to expire data in MongoDB Atlas?",
        "What is vector search index in Mongo Atlas?",
        "How does MongoDB Atlas different from relational DB in terms of data modeling",
    ],
    "Database Essentials": [
        "What is the impact of interleaving transactions in database operations?",
        "What is vector search index? how is it related to semantic search?",
    ],
}

test_set = { "MongoDB Atlas": [ "How do you secure MongoDB Atlas?", "How can Time to Live (TTL) be used to expire data in MongoDB Atlas?", "What is vector search index in Mongo Atlas?", "How does MongoDB Atlas different from relational DB in terms of data modeling", ], "Database Essentials": [ "What is the impact of interleaving transactions in database operations?", "What is vector search index? how is it related to semantic search?", ], }

Alternatively, we can generate test set automatically¶

In [ ]:

# test = GenerateTestSet(app_callable = query_engine.query)
# Generate the test set of a specified breadth and depth without examples automatically
from trulens.benchmark.generate.generate_test_set import GenerateTestSet
test = GenerateTestSet(app_callable=query_engine.query)
test_set_autogenerated = test.generate_test_set(test_breadth=3, test_depth=2)

# test = GenerateTestSet(app_callable = query_engine.query) # Generate the test set of a specified breadth and depth without examples automatically from trulens.benchmark.generate.generate_test_set import GenerateTestSet test = GenerateTestSet(app_callable=query_engine.query) test_set_autogenerated = test.generate_test_set(test_breadth=3, test_depth=2)

Get testing!¶

Our test set is made up of 2 topics (test breadth), each with 2-3 questions (test depth).

We can store the topic as record level metadata and then test queries from each topic, using tru_query_engine_recorder as a context manager.

In [ ]:

with tru_query_engine_recorder as recording:
    for category in test_set:
        recording.record_metadata = dict(prompt_category=category)
        test_prompts = test_set[category]
        for test_prompt in test_prompts:
            response = query_engine.query(test_prompt)

with tru_query_engine_recorder as recording: for category in test_set: recording.record_metadata = dict(prompt_category=category) test_prompts = test_set[category] for test_prompt in test_prompts: response = query_engine.query(test_prompt)

Check evaluation results¶

Evaluation results can be viewed in the TruLens dashboard (started at the top of the notebook) or directly in the notebook.

In [ ]:

session.get_leaderboard()

session.get_leaderboard()

Perhaps if we use metadata filters to create specialized query engines, we can improve the search results and thus, the overall evaluation results.

But it may be clunky to have two separate query engines - then we have to decide which one to use!

Instead, let's use a router query engine to choose the query engine based on the query.

Router Query Engine + Metadata Filters¶

In [ ]:

# Specify metadata filters
metadata_filters_db_essentials = MetadataFilters(
    filters=[
        ExactMatchFilter(key="metadata.file_name", value="DBEssential-2021.pdf")
    ]
)
metadata_filters_atlas = MetadataFilters(
    filters=[
        ExactMatchFilter(
            key="metadata.file_name", value="atlas_best_practices.pdf"
        )
    ]
)

metadata_filters_databrick = MetadataFilters(
    filters=[
        ExactMatchFilter(
            key="metadata.file_name", value="DataBrick_vector_search.pdf"
        )
    ]
)
# Instantiate Atlas Vector Search as a retriever for each set of filters
vector_store_retriever_db_essentials = VectorIndexRetriever(
    index=vector_store_index,
    filters=metadata_filters_db_essentials,
    similarity_top_k=5,
)
vector_store_retriever_atlas = VectorIndexRetriever(
    index=vector_store_index, filters=metadata_filters_atlas, similarity_top_k=5
)
vector_store_retriever_databrick = VectorIndexRetriever(
    index=vector_store_index,
    filters=metadata_filters_databrick,
    similarity_top_k=5,
)
# Pass the retrievers into the query engines
query_engine_with_filters_db_essentials = RetrieverQueryEngine(
    retriever=vector_store_retriever_db_essentials
)
query_engine_with_filters_atlas = RetrieverQueryEngine(
    retriever=vector_store_retriever_atlas
)
query_engine_with_filters_databrick = RetrieverQueryEngine(
    retriever=vector_store_retriever_databrick
)

# Specify metadata filters metadata_filters_db_essentials = MetadataFilters( filters=[ ExactMatchFilter(key="metadata.file_name", value="DBEssential-2021.pdf") ] ) metadata_filters_atlas = MetadataFilters( filters=[ ExactMatchFilter( key="metadata.file_name", value="atlas_best_practices.pdf" ) ] ) metadata_filters_databrick = MetadataFilters( filters=[ ExactMatchFilter( key="metadata.file_name", value="DataBrick_vector_search.pdf" ) ] ) # Instantiate Atlas Vector Search as a retriever for each set of filters vector_store_retriever_db_essentials = VectorIndexRetriever( index=vector_store_index, filters=metadata_filters_db_essentials, similarity_top_k=5, ) vector_store_retriever_atlas = VectorIndexRetriever( index=vector_store_index, filters=metadata_filters_atlas, similarity_top_k=5 ) vector_store_retriever_databrick = VectorIndexRetriever( index=vector_store_index, filters=metadata_filters_databrick, similarity_top_k=5, ) # Pass the retrievers into the query engines query_engine_with_filters_db_essentials = RetrieverQueryEngine( retriever=vector_store_retriever_db_essentials ) query_engine_with_filters_atlas = RetrieverQueryEngine( retriever=vector_store_retriever_atlas ) query_engine_with_filters_databrick = RetrieverQueryEngine( retriever=vector_store_retriever_databrick )

In [ ]:

from llama_index.core.tools import QueryEngineTool

# Set up the two distinct tools (query engines)

essentials_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine_with_filters_db_essentials,
    description=("Useful for retrieving context about database essentials"),
)

atlas_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine_with_filters_atlas,
    description=("Useful for retrieving context about MongoDB Atlas"),
)

databrick_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine_with_filters_databrick,
    description=(
        "Useful for retrieving context about Databrick's course on Vector Databases and Search"
    ),
)

from llama_index.core.tools import QueryEngineTool # Set up the two distinct tools (query engines) essentials_tool = QueryEngineTool.from_defaults( query_engine=query_engine_with_filters_db_essentials, description=("Useful for retrieving context about database essentials"), ) atlas_tool = QueryEngineTool.from_defaults( query_engine=query_engine_with_filters_atlas, description=("Useful for retrieving context about MongoDB Atlas"), ) databrick_tool = QueryEngineTool.from_defaults( query_engine=query_engine_with_filters_databrick, description=( "Useful for retrieving context about Databrick's course on Vector Databases and Search" ), )

In [ ]:

# Create the router query engine
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import PydanticSingleSelector

router_query_engine = RouterQueryEngine(
    selector=PydanticSingleSelector.from_defaults(),
    query_engine_tools=[essentials_tool, atlas_tool, databrick_tool],
)

# Create the router query engine from llama_index.core.query_engine import RouterQueryEngine from llama_index.core.selectors import PydanticSingleSelector router_query_engine = RouterQueryEngine( selector=PydanticSingleSelector.from_defaults(), query_engine_tools=[essentials_tool, atlas_tool, databrick_tool], )

In [ ]:

from trulens.apps.llamaindex import TruLlama

tru_query_engine_recorder_with_router = TruLlama(
    router_query_engine,
    app_name="RAG",
    app_version="Router Query Engine + Filters v2",
    feedbacks=[f_groundedness, f_answer_relevance, f_context_relevance],
)

from trulens.apps.llamaindex import TruLlama tru_query_engine_recorder_with_router = TruLlama( router_query_engine, app_name="RAG", app_version="Router Query Engine + Filters v2", feedbacks=[f_groundedness, f_answer_relevance, f_context_relevance], )

In [ ]:

with tru_query_engine_recorder_with_router as recording:
    for category in test_set:
        recording.record_metadata = dict(prompt_category=category)
        test_prompts = test_set[category]
        for test_prompt in test_prompts:
            response = router_query_engine.query(test_prompt)

with tru_query_engine_recorder_with_router as recording: for category in test_set: recording.record_metadata = dict(prompt_category=category) test_prompts = test_set[category] for test_prompt in test_prompts: response = router_query_engine.query(test_prompt)

Check results!¶

In [ ]:

session.get_leaderboard()

session.get_leaderboard()