Adding Semantic Caching and Memory to Your RAG Application Using MongoDB and LangChain

1
! pip install -qU datasets langchain langchain-mongodb langchain-openai pymongo pandas

1
import getpass
2
OPENAI_API_KEY = getpass.getpass("Enter your OpenAI API key:")

1
    {
2
      "fields": [
3
        {
4
          "numDimensions": 1536,
5
          "path": "embedding",
6
          "similarity": "cosine",
7
          "type": "vector"
8
        }
9
      ]
10
    }

1
    {
2
      "fields": [
3
        {
4
          "numDimensions": 1536,
5
          "path": "embedding",
6
          "similarity": "cosine",
7
          "type": "vector"
8
        },
9
        {
10
          "path": "llm_string",
11
          "type": "filter"
12
        }
13
      ]
14
    }

1
MONGODB_URI = getpass.getpass("Enter your MongoDB connection string:")

1
from datasets import load_dataset
2
import pandas as pd
3
4
data = load_dataset("MongoDB/embedded_movies")
5
df = pd.DataFrame(data["train"])
6
7
# Only keep records where the fullplot field is not null
8
df = df[df["fullplot"].notna()]
9
10
# Renaming the embedding field to "embedding" -- required by LangChain
11
df.rename(columns={"plot_embedding": "embedding"}, inplace=True)

1
from pymongo import MongoClient
2
3
# Initialize MongoDB python client
4
client = MongoClient(MONGODB_URI)
5
6
DB_NAME = "langchain_chatbot"
7
COLLECTION_NAME = "data"
8
ATLAS_VECTOR_SEARCH_INDEX_NAME = "vector_index"
9
collection = client[DB_NAME][COLLECTION_NAME]

1
    {
2
      "fields": [
3
        {
4
          "numDimensions": 1536,
5
          "path": "embedding",
6
          "similarity": "cosine",
7
          "type": "vector"
8
        }
9
      ]
10
    }

1
# Delete any existing records in the collection
2
collection.delete_many({})
3
4
# Data Ingestion
5
records = df.to_dict('records')
6
collection.insert_many(records)
7
8
print("Data ingestion into MongoDB completed")

1
from langchain_openai import OpenAIEmbeddings
2
from langchain_mongodb import MongoDBAtlasVectorSearch
3
4
# Using the text-embedding-ada-002 since that's what was used to create embeddings in the movies dataset
5
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY, model="text-embedding-ada-002")
6
7
# Vector Store Creation
8
vector_store = MongoDBAtlasVectorSearch.from_connection_string(
9
    connection_string=MONGODB_URI,
10
    namespace=DB_NAME + "." + COLLECTION_NAME,
11
    embedding= embeddings,
12
    index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
13
    text_key="fullplot"
14
)

1
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})

1
from langchain_openai import ChatOpenAI
2
from langchain_core.prompts import ChatPromptTemplate
3
from langchain_core.runnables import RunnablePassthrough
4
from langchain_core.output_parsers import StrOutputParser
5
6
# Generate context using the retriever, and pass the user question through
7
retrieve = {"context": retriever | (lambda docs: "\n\n".join([d.page_content for d in docs])), "question": RunnablePassthrough()}
8
template = """Answer the question based only on the following context: \
9
{context}
10
11
Question: {question}
12
"""
13
# Defining the chat prompt
14
prompt = ChatPromptTemplate.from_template(template)
15
# Defining the model to be used for chat completion
16
model = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
17
# Parse output as a string
18
parse_output = StrOutputParser()
19
20
# Naive RAG chain 
21
naive_rag_chain = (
22
    retrieve
23
    | prompt
24
    | model
25
    | parse_output
26
)

1
naive_rag_chain.invoke("What is the best movie to watch when sad?")
2
Output: Once a Thief

1
    export LANGCHAIN_TRACING_V2=true
2
    export LANGCHAIN_API_KEY=<your-api-key>

1
from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
2
from langchain_core.runnables.history import RunnableWithMessageHistory
3
from langchain_core.prompts import MessagesPlaceholder
4
5
def get_session_history(session_id: str) -> MongoDBChatMessageHistory:
6
        return MongoDBChatMessageHistory(MONGODB_URI, session_id, database_name=DB_NAME, collection_name="history")

1
standalone_system_prompt = """
2
Given a chat history and a follow-up question, rephrase the follow-up question to be a standalone question. \
3
Do NOT answer the question, just reformulate it if needed, otherwise return it as is. \
4
Only return the final standalone question. \
5
"""
6
standalone_question_prompt = ChatPromptTemplate.from_messages(
7
    [
8
        ("system", standalone_system_prompt),
9
        MessagesPlaceholder(variable_name="history"),
10
        ("human", "{question}"),
11
    ]
12
)
13
14
question_chain = standalone_question_prompt | model | parse_output

1
retriever_chain = RunnablePassthrough.assign(context=question_chain | retriever | (lambda docs: "\n\n".join([d.page_content for d in docs])))

1
rag_system_prompt = """Answer the question based only on the following context: \
2
{context}
3
"""
4
rag_prompt = ChatPromptTemplate.from_messages(
5
    [
6
        ("system", rag_system_prompt),
7
        MessagesPlaceholder(variable_name="history"),
8
        ("human", "{question}"),
9
    ]
10
)

1
# RAG chain
2
rag_chain = (
3
    retriever_chain
4
    | rag_prompt
5
    | model
6
    | parse_output
7
)
8
9
# RAG chain with history
10
with_message_history = RunnableWithMessageHistory(
11
    rag_chain,
12
    get_session_history,
13
    input_messages_key="question",
14
    history_messages_key="history",
15
)

1
with_message_history.invoke({"question": "What is the best movie to watch when sad?"}, {"configurable": {"session_id": "1"}})
2
3
Output: The best movie to watch when feeling down could be "Last Action Hero." It\'s a fun and action-packed film that blends reality and fantasy, offering an escape from the real world and providing an entertaining distraction.
4
5
with_message_history.invoke({"question": "How about something more light?"}, {"configurable": {"session_id": "1"}})
6
7
Output: For a lighter movie option, you might enjoy "Cousins." It\'s a comedy film set in Barcelona with action and humor, offering a fun and entertaining escape from reality. The storyline is engaging and filled with comedic moments that could help lift your spirits.

1
from langchain_mongodb.cache import MongoDBAtlasSemanticCache
2
from langchain_core.globals import set_llm_cache
3
4
set_llm_cache(MongoDBAtlasSemanticCache(
5
    connection_string=MONGODB_URI,
6
    embedding=embeddings,
7
    collection_name="semantic_cache",
8
    database_name=DB_NAME,
9
    index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
10
    wait_until_ready=True # Optional, waits until the cache is ready to be used
11
))

1
%%time
2
naive_rag_chain.invoke("What is the best movie to watch when sad?")

1
CPU times: user 87.8 ms, sys: 670 µs, total: 88.5 ms
2
Wall time: 1.24 s
3
'Once a Thief'

1
%%time
2
naive_rag_chain.invoke("Which movie do I watch when sad?")

1
CPU times: user 43.5 ms, sys: 4.16 ms, total: 47.7 ms
2
Wall time: 255 ms
3
'Once a Thief'