-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathrun.py
66 lines (60 loc) · 2.23 KB
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import os
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.memory import ConversationBufferMemory
from langchain.agents import Tool
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.tools import DuckDuckGoSearchRun
from langchain.callbacks import get_openai_callback
from langchain.schema import (
AIMessage,
HumanMessage,
SystemMessage
)
# authentication
with open('auth.txt', 'r') as f:
auth = f.readlines()[0].strip()
os.environ['OPENAI_API_KEY'] = auth
cost = 0
# set up GPT-3.5 Chat Bot
chat = ChatOpenAI(temperature=0)
history = []
# set up GPT-3.5 agent tools
llm = OpenAI()
agentTools = [DuckDuckGoSearchRun()]
agentTools.extend(load_tools(['llm-math', 'terminal'], llm=llm))
# set up GPT-3.5 director
dllm = ChatOpenAI(temperature=0)
directorSys = 'Does the following request require searching the internet, interacting with the filesystem, ' \
'executing code, or doing math calculations?\nRespond only with yes or no.'
def direct(query):
query = '"{}"'.format(query)
yn = dllm([SystemMessage(content=directorSys), HumanMessage(content=query)]).content.lower()
return 'yes' in yn
# Main loop
print('Type exit to quit')
while 1:
prompt = input('> ')
if prompt.lower() == 'exit':
print('Total cost: $' + str(cost))
exit()
with get_openai_callback() as cb:
history.append(HumanMessage(content=prompt))
answer = ''
if direct(prompt):
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
# give the agent 1 message of context if there is any
if len(history) > 1:
memory.chat_memory.add_ai_message(history[-2].content)
agent = initialize_agent(tools=agentTools, llm=chat, agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=True, memory=memory)
try:
answer = agent.run(prompt)
except Exception as e:
print('ERROR: ' + str(e))
else:
answer = chat(history).content
history.append(AIMessage(content=answer))
cost += cb.total_cost
print(answer)