Update rag.py

rag.py

@@ -99,9 +99,8 @@ def llm_chain(config, openai_api_key, prompt):
     llm_chain = LLMChain(llm = get_llm(config, openai_api_key), 
                          prompt = LLM_CHAIN_PROMPT)
     
-    completion = llm_chain.generate([{"question": prompt}])
-
     with get_openai_callback() as cb:
+        completion = llm_chain.generate([{"question": prompt}])
         print(cb)
     
     return completion, llm_chain
@@ -123,6 +122,4 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
         completion = rag_chain({"query": prompt})
         print(cb)
 
-    print("###" + str(completion))
-
     return completion, rag_chain