Update rag_langchain.py

rag_langchain.py

@@ -15,121 +15,124 @@ from langchain.vectorstores import MongoDBAtlasVectorSearch
 
 from pymongo import MongoClient
 
-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"
-YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
-
-CHROMA_DIR  = "/data/db"
-YOUTUBE_DIR = "/data/yt"
-
-MONGODB_ATLAS_CLUSTER_URI = os.environ["MONGODB_ATLAS_CLUSTER_URI"]
-MONGODB_DB_NAME           = "langchain_db"
-MONGODB_COLLECTION_NAME   = "gpt-4"
-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"])
-
-logging.basicConfig(stream = sys.stdout, level = logging.INFO)
-logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))
-
-def load_documents():
-    docs = []
+#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"
+#YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
+
+#MONGODB_ATLAS_CLUSTER_URI = os.environ["MONGODB_ATLAS_CLUSTER_URI"]
+#MONGODB_DB_NAME           = "langchain_db"
+#MONGODB_COLLECTION_NAME   = "gpt-4"
+#MONGODB_INDEX_NAME        = "default"
+
+#logging.basicConfig(stream = sys.stdout, level = logging.INFO)
+#logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))
+
+class LangChainRAG:
+    MONGODB_DB_NAME = "langchain_db"
+    
+    CHROMA_DIR  = "/data/db"
+    YOUTUBE_DIR = "/data/yt"
+
+    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"])
+
+    def load_documents(self):
+        docs = []
     
-    # PDF
-    loader = PyPDFLoader(PDF_URL)
-    docs.extend(loader.load())
-    #print("docs = " + str(len(docs)))
+        # PDF
+        loader = PyPDFLoader(self.PDF_URL)
+        docs.extend(loader.load())
+        #print("docs = " + str(len(docs)))
     
-    # Web
-    loader = WebBaseLoader(WEB_URL)
-    docs.extend(loader.load())
-    #print("docs = " + str(len(docs)))
+        # Web
+        loader = WebBaseLoader(self.WEB_URL)
+        docs.extend(loader.load())
+        #print("docs = " + str(len(docs)))
     
-    # YouTube
-    loader = GenericLoader(
-        YoutubeAudioLoader(
-            [YOUTUBE_URL_1, YOUTUBE_URL_2], 
-            YOUTUBE_DIR), 
-        OpenAIWhisperParser())
-    docs.extend(loader.load())
-    #print("docs = " + str(len(docs)))
+        # YouTube
+        loader = GenericLoader(
+            YoutubeAudioLoader(
+                [self.YOUTUBE_URL_1, self.YOUTUBE_URL_2], 
+                self.YOUTUBE_DIR), 
+            OpenAIWhisperParser())
+        docs.extend(loader.load())
+        #print("docs = " + str(len(docs)))
     
-    return docs
+        return docs
 
-def split_documents(config, docs):
-    text_splitter = RecursiveCharacterTextSplitter()
+    def split_documents(self, config, docs):
+        text_splitter = RecursiveCharacterTextSplitter()
     
-    return text_splitter.split_documents(docs)
+        return text_splitter.split_documents(docs)
     
-def store_documents_chroma(chunks):
-    Chroma.from_documents(
-        documents = chunks, 
-        embedding = OpenAIEmbeddings(disallowed_special = ()), 
-        persist_directory = CHROMA_DIR)
-
-def store_documents_mongodb(chunks):
-    client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
-    collection = client[MONGODB_DB_NAME][MONGODB_COLLECTION_NAME]
-
-    MongoDBAtlasVectorSearch.from_documents(
-        documents = chunks,
-        embedding = OpenAIEmbeddings(disallowed_special = ()),
-        collection = collection,
-        index_name = MONGODB_INDEX_NAME)
-
-def rag_ingestion_langchain(config):
-    docs = load_documents()
+    def store_documents_chroma(self, chunks):
+        Chroma.from_documents(
+            documents = chunks, 
+            embedding = OpenAIEmbeddings(disallowed_special = ()), 
+            persist_directory = self.CHROMA_DIR)
+
+    def store_documents_mongodb(self, chunks):
+        client = MongoClient(self.MONGODB_ATLAS_CLUSTER_URI)
+        collection = client[self.MONGODB_DB_NAME][self.MONGODB_COLLECTION_NAME]
+
+        MongoDBAtlasVectorSearch.from_documents(
+            documents = chunks,
+            embedding = OpenAIEmbeddings(disallowed_special = ()),
+            collection = collection,
+            index_name = self.MONGODB_INDEX_NAME)
+
+    def rag_ingestion_langchain(self, config):
+        docs = self.load_documents()
     
-    chunks = split_documents(config, docs)
+        chunks = self.split_documents(config, docs)
     
-    #store_documents_chroma(chunks)
-    store_documents_mongodb(chunks)
-
-def get_vector_store_chroma():
-    return Chroma(
-        embedding_function = OpenAIEmbeddings(disallowed_special = ()),
-        persist_directory = CHROMA_DIR)
-
-def get_vector_store_mongodb():
-    return MongoDBAtlasVectorSearch.from_connection_string(
-        MONGODB_ATLAS_CLUSTER_URI,
-        MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
-        OpenAIEmbeddings(disallowed_special = ()),
-        index_name = MONGODB_INDEX_NAME)
-
-def get_llm(config):
-    return ChatOpenAI(
-        model_name = config["model_name"], 
-        temperature = config["temperature"])
-
-def llm_chain(config, prompt):
-    llm_chain = LLMChain(
-        llm = get_llm(config), 
-        prompt = LLM_CHAIN_PROMPT)
+        #self.store_documents_chroma(chunks)
+        self.store_documents_mongodb(chunks)
+
+    def get_vector_store_chroma(self):
+        return Chroma(
+            embedding_function = OpenAIEmbeddings(disallowed_special = ()),
+            persist_directory = self.CHROMA_DIR)
+
+    def get_vector_store_mongodb(self):
+        return MongoDBAtlasVectorSearch.from_connection_string(
+            self.MONGODB_ATLAS_CLUSTER_URI,
+            self.MONGODB_DB_NAME + "." + self.MONGODB_COLLECTION_NAME,
+            OpenAIEmbeddings(disallowed_special = ()),
+            index_name = self.MONGODB_INDEX_NAME)
+
+    def get_llm(self, config):
+        return ChatOpenAI(
+            model_name = config["model_name"], 
+            temperature = config["temperature"])
+
+    def llm_chain(self, config, prompt):
+        llm_chain = LLMChain(
+            llm = self.get_llm(config), 
+            prompt = self.LLM_CHAIN_PROMPT)
     
-    with get_openai_callback() as callback:
-        completion = llm_chain.generate([{"question": prompt}])
+        with get_openai_callback() as callback:
+            completion = llm_chain.generate([{"question": prompt}])
     
-    return completion, llm_chain, callback
-
-def rag_chain(config, prompt):
-    #vector_store = get_vector_store_chroma()
-    vector_store = get_vector_store_mongodb()
-
-    rag_chain = RetrievalQA.from_chain_type(
-        get_llm(config), 
-        chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT,
-                             "verbose": True}, 
-        retriever = vector_store.as_retriever(search_kwargs = {"k": config["k"]}), 
-        return_source_documents = True)
+        return completion, llm_chain, callback
+
+    def rag_chain(self, config, prompt):
+        #vector_store = self.get_vector_store_chroma()
+        vector_store = self.get_vector_store_mongodb()
+
+        rag_chain = RetrievalQA.from_chain_type(
+            self.get_llm(config), 
+            chain_type_kwargs = {"prompt": self.RAG_CHAIN_PROMPT,
+                                 "verbose": True}, 
+            retriever = vector_store.as_retriever(search_kwargs = {"k": config["k"]}), 
+            return_source_documents = True)
     
-    with get_openai_callback() as callback:
-        completion = rag_chain({"query": prompt})
+        with get_openai_callback() as callback:
+            completion = rag_chain({"query": prompt})
 
-    return completion, rag_chain, callback
+        return completion, rag_chain, callback