Update rag.py

rag.py

@@ -1,5 +1,6 @@
 import openai, os
 
+from langchain.callbacks import get_openai_callback
 from langchain.chains import LLMChain, RetrievalQA
 from langchain.chat_models import ChatOpenAI
 from langchain.document_loaders import PyPDFLoader, WebBaseLoader
@@ -105,9 +106,6 @@ def llm_chain(config, openai_api_key, prompt):
 def rag_chain(config, openai_api_key, rag_option, prompt):
     llm = get_llm(config, openai_api_key)
 
-    with get_openai_callback() as cb:
-        print(cb)
-    
     if (rag_option == RAG_CHROMA):
         db = document_retrieval_chroma()
     elif (rag_option == RAG_MONGODB):
@@ -120,5 +118,8 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
     
     completion = rag_chain({"query": prompt})
     print("###" + str(completion))
+
+    with get_openai_callback() as cb:
+        print(cb)
     
     return completion, rag_chain