Update rag.py

rag.py

@@ -81,15 +81,23 @@ def run_rag_batch(config):
     embed_store_documents_mongodb(chunks)
 
 def retrieve_documents_chroma():
-    return Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
-                  persist_directory = CHROMA_DIR)
+    with get_openai_callback() as cb:
+        db = Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
+                    persist_directory = CHROMA_DIR)
+        print(cb)
+
+    return db
 
 def retrieve_documents_mongodb():
-    return MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
-                                                           MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
-                                                           OpenAIEmbeddings(disallowed_special = ()),
-                                                           index_name = MONGODB_INDEX_NAME)
+    with get_openai_callback() as cb:
+        db = MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
+                                                             MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
+                                                             OpenAIEmbeddings(disallowed_special = ()),
+                                                             index_name = MONGODB_INDEX_NAME)
+        print cb
 
+    return db
+    
 def get_llm(config, openai_api_key):
     return ChatOpenAI(model_name = config["model_name"], 
                       openai_api_key = openai_api_key,