Update rag.py

rag.py

@@ -103,7 +103,7 @@ def llm_chain(config, openai_api_key, prompt):
         completion = llm_chain.generate([{"question": prompt}])
         print(cb)
     
-    return completion, llm_chain
+    return completion, llm_chain, cb
 
 def rag_chain(config, openai_api_key, rag_option, prompt):
     llm = get_llm(config, openai_api_key)
@@ -122,4 +122,4 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
         completion = rag_chain({"query": prompt})
         print(cb)
 
-    return completion, rag_chain
+    return completion, rag_chain, cb