Update agent_llamaindex.py

agent_llamaindex.py

@@ -1,10 +1,8 @@
-import os, tiktoken
+import os
 
 from datetime import date
 from llama_hub.tools.weather import OpenWeatherMapToolSpec
-from llama_index import ServiceContext
 from llama_index.agent import OpenAIAgent
-from llama_index.callbacks import CallbackManager, TokenCountingHandler
 from llama_index.llms import OpenAI
 from llama_index.tools import FunctionTool
 
@@ -14,54 +12,21 @@ def today_tool(text: str) -> str:
        Any date mathematics should occur outside this function."""
     return str(date.today())
 
-def get_callback_manager(config):
-    token_counter = TokenCountingHandler(
-        tokenizer = tiktoken.encoding_for_model(config["model"]).encode
-    )
-
-    token_counter.reset_counts()
-
-    return CallbackManager([token_counter])
-
-def get_callback(token_counter):
-    return ("Tokens Used: " +
-            str(token_counter.total_llm_token_count) + "\n" +
-            "Prompt Tokens: " +
-            str(token_counter.prompt_llm_token_count) + "\n" +
-            "Completion Tokens: " +
-            str(token_counter.completion_llm_token_count))
-
-def get_llm(config):
-    return OpenAI(
-        model = config["model"], 
-        temperature = config["temperature"]
-    )
-        
-def get_service_context(config):
-    return ServiceContext.from_defaults(
-        callback_manager = get_callback_manager(config),
-        llm = get_llm(config)
-    )
-        
 def agent_llamaindex(config, prompt):
-    service_context = get_service_context(config)
-    
+    llm = OpenAI(
+        model = config["model"],
+        temperature = config["temperature"])
+
     tool_spec = OpenWeatherMapToolSpec(key = os.environ["OPENWEATHERMAP_API_KEY"])
     tools = tool_spec.to_tool_list()
     
     date_tool = FunctionTool.from_defaults(fn = today_tool)
-    
+            
     agent = OpenAIAgent.from_tools(
         [tools[0],   # built-in tools
          date_tool], # custom tools
-        llm = get_llm(config),
-        service_context = service_context,
+        llm = llm, 
         verbose = True
     )
 
-    completion = agent.chat(prompt)
-
-    callback = get_callback(
-        service_context.callback_manager.handlers[0])
-
-    return completion, callback
+    return agent.chat(prompt)