Update rag.py

rag.py

@@ -98,14 +98,22 @@ def llm_chain(openai_api_key, prompt):
     llm_chain = LLMChain(llm = get_llm(openai_api_key), 
                          prompt = LLM_CHAIN_PROMPT, 
                          verbose = False)
+    
     completion = llm_chain.generate([{"question": prompt}])
+    
     return completion, llm_chain
 
-def rag_chain(openai_api_key, prompt, db):
-    rag_chain = RetrievalQA.from_chain_type(get_llm(openai_api_key), 
+def rag_chain(openai_api_key, prompt):
+    llm = get_llm(openai_api_key)
+    
+    db = document_retrieval_chroma(llm, prompt)
+    
+    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