bstraehle commited on
Commit
82394a2
·
1 Parent(s): e472bde

Update agent_langchain.py

Browse files
Files changed (1) hide show
  1. agent_langchain.py +7 -1
agent_langchain.py CHANGED
@@ -2,6 +2,7 @@ import os
2
 
3
  from datetime import date
4
  from langchain.agents import AgentType, initialize_agent, load_tools, tool
 
5
  from langchain.chat_models import ChatOpenAI
6
 
7
  os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
@@ -33,4 +34,9 @@ def agent_langchain(config, prompt):
33
  verbose = True
34
  )
35
 
36
- return agent(prompt)
 
 
 
 
 
 
2
 
3
  from datetime import date
4
  from langchain.agents import AgentType, initialize_agent, load_tools, tool
5
+ from langchain.callbacks import get_openai_callback
6
  from langchain.chat_models import ChatOpenAI
7
 
8
  os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
 
34
  verbose = True
35
  )
36
 
37
+ with get_openai_callback() as callback:
38
+ completion = agent(prompt)
39
+
40
+ print(str(callback))
41
+
42
+ return completion