Add Memory and Semantic Caching with LangChain and MongoDB

Set up the environment.

Set up the environment for this tutorial. Create an interactive Python notebook by saving a file with the .ipynb extension. This notebook allows you to run Python code snippets individually, and you'll use it to run the code in this tutorial.

To set up your notebook environment:

1. Run the following command in your notebook:

   pip install --quiet --upgrade langchain langchain-community langchain-core langchain-mongodb langchain-voyageai langchain-openai langgraph langgraph-checkpoint-mongodb pypdf

2. Set environment variables.

   Run the following code to set the environment variables for this tutorial. Provide your Voyage API key, OpenAI API Key, and MongoDB cluster's SRV connection string.

   import os
   os.environ["OPENAI_API_KEY"] = "<openai-key>"
   os.environ["VOYAGE_API_KEY"] = "<voyage-key>"
   MONGODB_URI = "<connection-string>"

   Note

   Replace <connection-string> with the connection string for your Atlas cluster or local Atlas deployment.

   Atlas Cluster

   Local or Self-Managed

   Your connection string should use the following format:

   mongodb+srv://<db_username>:<db_password>@<clusterName>.<hostname>.mongodb.net

   To learn more, see Connect to a Cluster via Client Libraries.

   Your connection string should use the following format:

   mongodb://localhost:<port-number>/?directConnection=true

   To learn more, see Connection Strings.

Instantiate the vector store.

Paste and run the following code in your notebook to create a vector store instance named vector_store using the langchain_db.rag_with_memory namespace in MongoDB:

from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_voyageai import VoyageAIEmbeddings
# Use the voyage-3-large embedding model
embedding_model = VoyageAIEmbeddings(model="voyage-3-large")
# Create the vector store
vector_store = MongoDBAtlasVectorSearch.from_connection_string(
   connection_string = MONGODB_URI,
   embedding = embedding_model,
   namespace = "langchain_db.rag_with_memory"
)

Add data to the vector store.

Paste and run the following code in your notebook to ingest a sample PDF that contains a recent MongoDB earnings report into the vector store.

This code uses a text splitter to chunk the PDF data into smaller parent documents. It specifies the chunk size (number of characters) and chunk overlap (number of overlapping characters between consecutive chunks) for each document.

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
# Load the PDF
loader = PyPDFLoader("https://investors.mongodb.com/node/13176/pdf")
data = loader.load()
# Split PDF into documents
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
docs = text_splitter.split_documents(data)
# Add data to the vector store
vector_store.add_documents(docs)

Create the MongoDB Vector Search index.

Run the following code to create the MongoDB Vector Search index for the vector store to enable vector search over your data:

# Use LangChain helper method to create the vector search index
vector_store.create_vector_search_index(
   dimensions = 1024 # The dimensions of the vector embeddings to be indexed
)

The index should take about one minute to build. While it builds, the index is in an initial sync state. When it finishes building, you can start querying the data in your collection.

Configure the MongoDB checkpointer.

To maintain conversation history across multiple interactions, use the LangGraph MongoDBSaver checkpointer. A checkpointer saves the state of a conversation to MongoDB after each step and restores it on the next interaction, which lets your RAG application handle conversation context.

LangGraph identifies each conversation by a thread_id. The checkpointer writes state to the checkpoints and checkpoint_writes collections in the database that you specify.

Paste and run the following code in your notebook to create a checkpointer named checkpointer that stores conversation state in the langchain_db database.

from langgraph.checkpoint.mongodb import MongoDBSaver
from pymongo import MongoClient
# Connect to your MongoDB cluster
mongo_client = MongoClient(MONGODB_URI)
# Create the checkpointer to persist conversation state
checkpointer = MongoDBSaver(mongo_client, db_name="langchain_db")

Create a RAG chain that handles chat message history.

Paste and run the following code snippets to create the RAG chain:

1. Specify the LLM to use.

   from langchain_openai import ChatOpenAI
   # Define the model to use for chat completion
   llm = ChatOpenAI(model = "gpt-4o")

2. Define a prompt that summarizes the chat history for the retriever.

   from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
   from langchain_core.output_parsers import StrOutputParser
   # Create a prompt to generate standalone questions from follow-up questions
   standalone_system_prompt = """
   Given a chat history and a follow-up question, rephrase the follow-up question to be a standalone question.
   Do NOT answer the question, just reformulate it if needed, otherwise return it as is.
   Only return the final standalone question.
   """
   standalone_question_prompt = ChatPromptTemplate.from_messages(
       [
           ("system", standalone_system_prompt),
           MessagesPlaceholder(variable_name="history"),
           ("human", "{question}"),
       ]
   )
   # Parse output as a string
   parse_output = StrOutputParser()
   question_chain = standalone_question_prompt | llm | parse_output

3. Build a retriever chain that processes the chat history and retrieves documents.

   from langchain_core.runnables import RunnablePassthrough
   # Create a retriever
   retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={ "k": 5 })
   # Create a retriever chain that processes the question with history and retrieves documents
   retriever_chain = RunnablePassthrough.assign(context=question_chain | retriever | (lambda docs: "\n\n".join([d.page_content for d in docs])))

4. Define a prompt to generate an answer based on the chat history and retrieved context.

   # Create a prompt template that includes the retrieved context and chat history
   rag_system_prompt = """Answer the question based only on the following context:
   {context}
   """
   rag_prompt = ChatPromptTemplate.from_messages(
       [
           ("system", rag_system_prompt),
           MessagesPlaceholder(variable_name="history"),
           ("human", "{question}"),
       ]
   )

5. Implement RAG with memory.

   Combine the components you defined into a complete RAG chain:

   # Build the RAG chain
   rag_chain = (
       retriever_chain
       | rag_prompt
       | llm
       | parse_output
   )

   Then, wrap the chain in a LangGraph graph and compile it with the checkpointer. The graph contains a single node that reads the conversation state, passes the latest question and the prior messages to the chain, and appends the answer to the state:

   from langgraph.graph import StateGraph, START, END, MessagesState
   from langchain_core.messages import AIMessage
   # Define the node that runs the RAG chain
   def call_rag_chain(state: MessagesState):
       question = state["messages"][-1].content
       history = state["messages"][:-1]
       answer = rag_chain.invoke(
           {"question": question, "history": history}
       )
       return {"messages": [AIMessage(content=answer)]}
   # Build the graph
   workflow = StateGraph(MessagesState)
   workflow.add_node("rag", call_rag_chain)
   workflow.add_edge(START, "rag")
   workflow.add_edge("rag", END)
   # Compile the graph with the checkpointer
   rag_with_memory = workflow.compile(checkpointer=checkpointer)

Test your RAG implementation.

Invoke the chain to answer questions. This chain maintains the conversation context and returns relevant answers that consider the previous interactions. Your responses might vary.

# First question
response_1 = rag_with_memory.invoke(
    {"messages": [("human", "What was MongoDB's latest acquisition?")]},
    {"configurable": {"thread_id": "user_1"}}
)
print(response_1["messages"][-1].content)

MongoDB's latest acquisition was Voyage AI, a pioneer in state-of-the-art embedding and reranking models for next-generation AI applications.

# Follow-up question that references the previous question
response_2 = rag_with_memory.invoke(
    {"messages": [("human", "Why did they do it?")]},
    {"configurable": {"thread_id": "user_1"}}
)
print(response_2["messages"][-1].content)

MongoDB acquired Voyage AI to enable organizations to easily build trustworthy AI applications by integrating advanced embedding and reranking models into their technology. This acquisition aligns with MongoDB's goal of helping businesses innovate at "AI speed" using its flexible document model and seamless scalability.

Configure the semantic cache.

Run the following code to configure the semantic cache by using the MongoDBAtlasSemanticCache class:

from langchain_mongodb.cache import MongoDBAtlasSemanticCache
from langchain_core.globals import set_llm_cache
# Configure the semantic cache
set_llm_cache(MongoDBAtlasSemanticCache(
    connection_string = MONGODB_URI,
    database_name = "langchain_db",
    collection_name = "semantic_cache",
    embedding = embedding_model,
    index_name = "vector_index",
    similarity_threshold = 0.5  # Adjust based on your requirements
))

Test the semantic cache with your RAG chain.

The semantic cache automatically caches your prompts. Run the following sample queries, where you should see a significant reduction in response time for the second query. Your responses and response times might vary.

%%time
# First query (not cached)
rag_with_memory.invoke(
  {"messages": [("human", "What was MongoDB's latest acquisition?")]},
  {"configurable": {"thread_id": "user_2"}}
)

CPU times: user 54.7 ms, sys: 34.2 ms, total: 88.9 ms
Wall time: 7.42 s
"MongoDB's latest acquisition was Voyage AI, a pioneer in state-of-the-art embedding and reranking models that power next-generation AI applications."

%%time
# Second query (cached)
rag_with_memory.invoke(
  {"messages": [("human", "What company did MongoDB acquire recently?")]},
  {"configurable": {"thread_id": "user_2"}}
)

CPU times: user 79.7 ms, sys: 24 ms, total: 104 ms
Wall time: 3.87 s
'MongoDB recently acquired Voyage AI.'