Update rag_llamaindex.py

rag_llamaindex.py

@@ -60,6 +60,20 @@ class LlamaIndexRAG(BaseRAG):
         token_counter.reset_counts()
 
         return CallbackManager([token_counter])
+
+    def get_callback(self, token_counter):
+        return "Embedding Tokens: ",
+                token_counter.total_embedding_token_count,
+                "\n",
+                "LLM Prompt Tokens: ",
+                token_counter.prompt_llm_token_count,
+                "\n",
+                "LLM Completion Tokens: ",
+                token_counter.completion_llm_token_count,
+                "\n",
+                "Total LLM Token Count: ",
+                token_counter.total_llm_token_count,
+                "\n"
         
     def get_llm(self, config):
         return OpenAI(
@@ -99,7 +113,7 @@ class LlamaIndexRAG(BaseRAG):
             service_context = self.get_service_context(config),
             storage_context = self.get_storage_context()
         )
-
+  
     def ingestion(self, config):
         docs = self.load_documents()
     
@@ -119,11 +133,7 @@ class LlamaIndexRAG(BaseRAG):
         )
 
         completion = query_engine.query(prompt)
-
-        print("111 " + str(service_context.callback_manager))
-        print("111 " + str(service_context.callback_manager.handlers))
-        print("111 " + str(service_context.callback_manager.handlers[0]))
-        token_counter = service_context.callback_manager.handlers[0]
-        print("222 " + str(token_counter.total_llm_token_count))
+        callback = get_callback(
+            service_context.callback_manager.handlers[0])
         
-        return completion
+        return completion, callback