Update rag.py

rag.py

@@ -1,4 +1,4 @@
-import os
+import openai, os
 
 from langchain.chains import LLMChain, RetrievalQA
 from langchain.chat_models import ChatOpenAI
@@ -38,6 +38,8 @@ config = {
     "chunk_overlap": 150,
     "chunk_size": 1500,
     "k": 3,
+    "model_name": "gpt-4-0613",
+    "temperature": 0,
 }
 
 def document_loading_splitting():
@@ -87,23 +89,23 @@ def document_retrieval_mongodb(llm, prompt):
                                                            OpenAIEmbeddings(disallowed_special = ()),
                                                            index_name = MONGODB_INDEX_NAME)
 
-def llm_chain(llm, prompt):
-    llm_chain = LLMChain(llm = llm, 
+def get_llm():
+    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(), 
                          prompt = LLM_CHAIN_PROMPT, 
                          verbose = False)
     completion = llm_chain.generate([{"question": prompt}])
     return completion, llm_chain
 
-def rag_chain(llm, prompt, db):
-    rag_chain = RetrievalQA.from_chain_type(llm, 
+def rag_chain(prompt, db):
+    rag_chain = RetrievalQA.from_chain_type(get_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
-
-def get_llm():
-    return ChatOpenAI(model_name = config["model_name"], 
-                      openai_api_key = openai_api_key, 
-                      temperature = config["temperature"])
+    return completion, rag_chain