Update rag.py

rag.py

@@ -1,3 +1,5 @@
+import os
+
 PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
 WEB_URL       = "https://openai.com/research/gpt-4"
 YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
@@ -14,3 +16,81 @@ MONGODB_INDEX_NAME        = "default"
 
 LLM_CHAIN_PROMPT = PromptTemplate(input_variables = ["question"], template = os.environ["LLM_TEMPLATE"])
 RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], template = os.environ["RAG_TEMPLATE"])
+
+RAG_OFF     = "Off"
+RAG_CHROMA  = "Chroma"
+RAG_MONGODB = "MongoDB"
+
+client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
+collection = client[MONGODB_DB_NAME][MONGODB_COLLECTION_NAME]
+
+config = {
+    "chunk_overlap": 150,
+    "chunk_size": 1500,
+    "k": 3,
+    "model_name": "gpt-4-0613",
+    "temperature": 0,
+}
+
+def document_loading_splitting():
+    # Document loading
+    docs = []
+    
+    # Load PDF
+    loader = PyPDFLoader(PDF_URL)
+    docs.extend(loader.load())
+    
+    # Load Web
+    loader = WebBaseLoader(WEB_URL)
+    docs.extend(loader.load())
+    
+    # Load YouTube
+    loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_1,
+                                               YOUTUBE_URL_2,
+                                               YOUTUBE_URL_3], YOUTUBE_DIR), 
+                           OpenAIWhisperParser())
+    docs.extend(loader.load())
+
+    # Document splitting
+    text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = config["chunk_overlap"],
+                                                   chunk_size = config["chunk_size"])
+    split_documents = text_splitter.split_documents(docs)
+    
+    return split_documents
+
+def document_storage_chroma(documents):
+    Chroma.from_documents(documents = documents, 
+                          embedding = OpenAIEmbeddings(disallowed_special = ()), 
+                          persist_directory = CHROMA_DIR)
+
+def document_storage_mongodb(documents):
+    MongoDBAtlasVectorSearch.from_documents(documents = documents,
+                                            embedding = OpenAIEmbeddings(disallowed_special = ()),
+                                            collection = collection,
+                                            index_name = MONGODB_INDEX_NAME)
+
+def document_retrieval_chroma(llm, prompt):
+    return Chroma(embedding_function = OpenAIEmbeddings(),
+                  persist_directory = CHROMA_DIR)
+
+def document_retrieval_mongodb(llm, prompt):
+    return MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
+                                                           MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
+                                                           OpenAIEmbeddings(disallowed_special = ()),
+                                                           index_name = MONGODB_INDEX_NAME)
+
+def llm_chain(llm, prompt):
+    llm_chain = LLMChain(llm = llm, 
+                         prompt = LLM_CHAIN_PROMPT, 
+                         verbose = False)
+    completion = llm_chain.generate([{"question": prompt}])
+    return completion, llm_chain
+
+def rag_chain(llm, prompt, db):
+    rag_chain = RetrievalQA.from_chain_type(llm, 
+                                            chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
+                                            retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
+                                            return_source_documents = True,
+                                            verbose = False)
+    completion = rag_chain({"query": prompt})
+    return completion, rag_chain