Update rag.py

rag.py

@@ -79,11 +79,11 @@ def rag_batch(config):
     document_storage_chroma(chunks)
     document_storage_mongodb(chunks)
 
-def document_retrieval_chroma():
+def document_retrieval_chroma(llm):
     return Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
                   persist_directory = CHROMA_DIR)
 
-def document_retrieval_mongodb():
+def document_retrieval_mongodb(llm):
     return MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
                                                            MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
                                                            OpenAIEmbeddings(disallowed_special = ()),
@@ -107,26 +107,17 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
     llm = get_llm(config, openai_api_key)
     
     if (rag_option == RAG_CHROMA):
-        db = document_retrieval_chroma()
+        db = document_retrieval_chroma(llm)
     elif (rag_option == RAG_MONGODB):
-        db = document_retrieval_mongodb()
-
-    ###
-    retriever = db.as_retriever(search_kwargs = {"k": config["k"]})
-    retrieved_docs = retriever.invoke(prompt)
-    print(retrieved_docs[0].page_content)
-    print(retrieved_docs[1].page_content)
-    print(retrieved_docs[2].page_content)
-    ###
-    
+        db = document_retrieval_mongodb(llm)
+
     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}, include_run_info = True)
+    completion = rag_chain({"query": prompt})
     print("###" + str(completion))
-    print("###" + str(completion["__run"]))
     
     return completion, rag_chain