Update agent_langchain.py

agent_langchain.py

@@ -2,6 +2,7 @@ import os
 
 from datetime import date
 from langchain.agents import AgentType, initialize_agent, load_tools, tool
+from langchain.callbacks import get_openai_callback
 from langchain.chat_models import ChatOpenAI
 
 os.environ["LANGCHAIN_ENDPOINT"]   = "https://api.smith.langchain.com"
@@ -23,14 +24,17 @@ def agent_langchain(config, prompt):
     OPENWEATHERMAP_API_KEY = os.environ["OPENWEATHERMAP_API_KEY"]
     
     tools = load_tools(["openweathermap-api"])
-            
-    agent = initialize_agent(
-        tools +      # built-in tools
-        [today_tool], # custom tools
-        llm,
-        agent = AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
-        handle_parsing_errors = True,
-        verbose = True
-    )
 
+    with get_openai_callback() as callback:
+        agent = initialize_agent(
+            tools +      # built-in tools
+            [today_tool], # custom tools
+            llm,
+            agent = AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+            handle_parsing_errors = True,
+            verbose = True
+        )
+
+    print(str(callback))
+    
     return agent(prompt)