Update rag.py

rag.py

@@ -81,13 +81,13 @@ def document_retrieval_mongodb(llm, prompt):
                                                            OpenAIEmbeddings(disallowed_special = ()),
                                                            index_name = MONGODB_INDEX_NAME)
 
-def get_llm(openai_api_key):
+def get_llm(config, openai_api_key):
     return ChatOpenAI(model_name = config["model_name"], 
                       openai_api_key = openai_api_key, 
                       temperature = config["temperature"])
 
-def llm_chain(openai_api_key, prompt):
-    llm_chain = LLMChain(llm = get_llm(openai_api_key), 
+def llm_chain(config, openai_api_key, prompt):
+    llm_chain = LLMChain(llm = get_llm(config, openai_api_key), 
                          prompt = LLM_CHAIN_PROMPT, 
                          verbose = False)
     
@@ -95,8 +95,8 @@ def llm_chain(openai_api_key, prompt):
     
     return completion, llm_chain
 
-def rag_chain(openai_api_key, prompt):
-    llm = get_llm(openai_api_key)
+def rag_chain(config, openai_api_key, prompt):
+    llm = get_llm(config, openai_api_key)
     
     db = document_retrieval_chroma(llm, prompt)