RAG template for MongoDB Atlas Vector Search (#12526)

templates/rag-mongo/rag_mongo/__init__.py

@@ -0,0 +1,3 @@
+from rag_mongo.chain import chain
+
+__all__ = ["chain"]

templates/rag-mongo/rag_mongo/chain.py

@@ -0,0 +1,79 @@
+import os
+
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel
+from langchain.schema.output_parser import StrOutputParser
+from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
+from langchain.vectorstores import MongoDBAtlasVectorSearch
+from pymongo import MongoClient
+
+# Set DB
+if os.environ.get("MONGO_URI", None) is None:
+    raise Exception("Missing `MONGO_URI` environment variable.")
+MONGO_URI = os.environ["MONGO_URI"]
+
+DB_NAME = "langchain-test-2"
+COLLECTION_NAME = "test"
+ATLAS_VECTOR_SEARCH_INDEX_NAME = "default"
+
+client = MongoClient(MONGO_URI)
+db = client[DB_NAME]
+MONGODB_COLLECTION = db[COLLECTION_NAME]
+
+### Ingest code - you may need to run this the first time
+""" 
+# Load
+from langchain.document_loaders import WebBaseLoader
+loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
+data = loader.load()
+
+# Split
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
+all_splits = text_splitter.split_documents(data)
+
+# Add to vectorDB
+# Insert the documents in MongoDB Atlas Vector Search
+vectorstore = MongoDBAtlasVectorSearch.from_documents(
+     documents=all_splits, 
+     embedding=OpenAIEmbeddings(disallowed_special=()), 
+     collection=MONGODB_COLLECTION, 
+     index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME
+ )
+retriever = vectorstore.as_retriever()
+"""
+
+# Read from MongoDB Atlas Vector Search
+vectorstore = MongoDBAtlasVectorSearch.from_connection_string(
+    MONGO_URI,
+    DB_NAME + "." + COLLECTION_NAME,
+    OpenAIEmbeddings(disallowed_special=()),
+    index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
+)
+retriever = vectorstore.as_retriever()
+
+# RAG prompt
+template = """Answer the question based only on the following context:
+{context}
+Question: {question}
+"""
+prompt = ChatPromptTemplate.from_template(template)
+
+# RAG
+model = ChatOpenAI()
+chain = (
+    RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
+    | prompt
+    | model
+    | StrOutputParser()
+)
+
+
+# Add typing for input
+class Question(BaseModel):
+    __root__: str
+
+
+chain = chain.with_types(input_type=Question)