Update rag.py

rag.py

@@ -89,20 +89,20 @@ def document_retrieval_mongodb(llm, prompt):
                                                            OpenAIEmbeddings(disallowed_special = ()),
                                                            index_name = MONGODB_INDEX_NAME)
 
-def get_llm():
+def get_llm(openai_api_key):
     return ChatOpenAI(model_name = config["model_name"], 
                       openai_api_key = openai_api_key, 
                       temperature = config["temperature"])
 
-def llm_chain(prompt):
-    llm_chain = LLMChain(llm = get_llm(), 
+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(prompt, db):
-    rag_chain = RetrievalQA.from_chain_type(get_llm(), 
+def rag_chain(openai_api_key, prompt, db):
+    rag_chain = RetrievalQA.from_chain_type(get_llm(openai_api_key), 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
                                             retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
                                             return_source_documents = True,