Update rag_langchain.py

rag_langchain.py

@@ -67,13 +67,13 @@ def split_documents(config, docs):
     
     return text_splitter.split_documents(docs)
     
-def store_chroma(chunks):
+def store_documents_chroma(chunks):
     Chroma.from_documents(
         documents = chunks, 
         embedding = OpenAIEmbeddings(disallowed_special = ()), 
         persist_directory = CHROMA_DIR)
 
-def store_mongodb(chunks):
+def store_documents_mongodb(chunks):
     client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
     collection = client[MONGODB_DB_NAME][MONGODB_COLLECTION_NAME]
 
@@ -88,15 +88,15 @@ def rag_ingestion(config):
     
     chunks = split_documents(config, docs)
     
-    #store_chroma(chunks)
-    store_mongodb(chunks)
+    #store_documents_chroma(chunks)
+    store_documents_mongodb(chunks)
 
-def retrieve_chroma():
+def get_vector_store_chroma():
     return Chroma(
         embedding_function = OpenAIEmbeddings(disallowed_special = ()),
         persist_directory = CHROMA_DIR)
 
-def retrieve_mongodb():
+def get_vector_store_mongodb():
     return MongoDBAtlasVectorSearch.from_connection_string(
         MONGODB_ATLAS_CLUSTER_URI,
         MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
@@ -113,23 +113,23 @@ def llm_chain(config, prompt):
         llm = get_llm(config), 
         prompt = LLM_CHAIN_PROMPT)
     
-    with get_openai_callback() as cb:
+    with get_openai_callback() as callback:
         completion = llm_chain.generate([{"question": prompt}])
     
-    return completion, llm_chain, cb
+    return completion, llm_chain, callback
 
 def rag_chain(config, prompt):
-    #db = retrieve_chroma()
-    db = retrieve_mongodb()
+    #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 = db.as_retriever(search_kwargs = {"k": config["k"]}), 
+        retriever = vector_store.as_retriever(search_kwargs = {"k": config["k"]}), 
         return_source_documents = True)
     
-    with get_openai_callback() as cb:
+    with get_openai_callback() as callback:
         completion = rag_chain({"query": prompt})
 
-    return completion, rag_chain, cb
+    return completion, rag_chain, callback