Update rag.py

rag.py

@@ -71,6 +71,14 @@ def document_storage_mongodb(chunks):
                                             collection = collection,
                                             index_name = MONGODB_INDEX_NAME)
 
+def rag_batch(config):
+    docs = document_loading()
+    
+    chunks = document_splitting(config, docs)
+    
+    document_storage_chroma(chunks)
+    document_storage_mongodb(chunks)
+
 def document_retrieval_chroma():
     return Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
                   persist_directory = CHROMA_DIR)
@@ -81,14 +89,6 @@ def document_retrieval_mongodb():
                                                            OpenAIEmbeddings(disallowed_special = ()),
                                                            index_name = MONGODB_INDEX_NAME)
 
-def rag_batch(config):
-    docs = document_loading()
-    
-    chunks = document_splitting(config, docs)
-    
-    document_storage_chroma(chunks)
-    document_storage_mongodb(chunks)
-
 def get_llm(config, openai_api_key):
     return ChatOpenAI(model_name = config["model_name"], 
                       openai_api_key = openai_api_key, 
@@ -110,6 +110,14 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
         db = document_retrieval_chroma()
     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)
+    ###
     
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
@@ -118,6 +126,7 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
                                             verbose = False)
     
     completion = rag_chain({"query": prompt}, include_run_info = True)
+    print("###" + str(completion))
     print("###" + str(completion["__run"]))
     
     return completion, rag_chain