Update rag.py

rag.py

@@ -119,10 +119,10 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
                                             retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
                                             return_source_documents = True)
     
-    completion = rag_chain({"query": prompt})
-    print("###" + str(completion))
-
     with get_openai_callback() as cb:
+        completion = rag_chain({"query": prompt})
         print(cb)
-    
+
+    print("###" + str(completion))
+
     return completion, rag_chain