138 lines
6.3 KiB
Python
138 lines
6.3 KiB
Python
#!/usr/bin/env python3
|
|
import json
|
|
import re
|
|
import readline
|
|
import signal
|
|
import socket
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
|
|
from openai import OpenAI
|
|
from termcolor import colored
|
|
|
|
from config import OPENAI_KEY
|
|
from lib.jsonify import jsonify_anything
|
|
from lib.openai.bash import func_run_bash
|
|
from lib.openai.functs import function_description, VALID_FUNCS
|
|
from lib.openai.google import search_google, search_google_maps, search_google_news
|
|
from lib.personality import load_personality
|
|
|
|
|
|
def signal_handler(sig, frame):
|
|
print()
|
|
sys.exit(0)
|
|
|
|
|
|
# Keep pycharm from removing this import.
|
|
readline.get_completion_type()
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
|
|
client = OpenAI(api_key=OPENAI_KEY)
|
|
|
|
bot_name = 'Sakura'
|
|
character_card = load_personality(bot_name, 'a shy girl', 'a desktop computer', 'she', 'Use Japanese emoticons.')
|
|
|
|
context: list[dict[str, str]] = [character_card]
|
|
|
|
|
|
def main():
|
|
print(colored(f'System Management Intelligence Interface', 'green', attrs=['bold']) + ' ' + colored(bot_name, 'green', attrs=['bold', 'underline']) + colored(' on ', 'green', attrs=['bold']) + colored(socket.gethostname(), 'green', attrs=['bold', 'underline']) + '\n')
|
|
|
|
while True:
|
|
try:
|
|
next_input = str(input('> '))
|
|
except EOFError:
|
|
print('Exit')
|
|
sys.exit(0)
|
|
print('')
|
|
context.append({'role': 'user', 'content': next_input})
|
|
|
|
i = 0
|
|
while True:
|
|
temp_context = context
|
|
if i > 0:
|
|
# Insert a prompt if this is not the first message.
|
|
temp_context.append(
|
|
{
|
|
'role': 'system',
|
|
'content': f"""Evaluate your progress on the current task. You have preformed {i} steps for this task so far. Use "end_my_response" if you are finished and ready for the user's response. Run another command using `run_bash` if necessary.
|
|
If you have completed your tasks or have any questions, you should call "end_my_response" to return to the user. The current time is {datetime.now()} {time.tzname[0]}."""}
|
|
)
|
|
|
|
response = client.chat.completions.create(
|
|
model="gpt-4-1106-preview", # TODO: config
|
|
messages=temp_context,
|
|
functions=function_description,
|
|
temperature=0.7 # TODO: config
|
|
)
|
|
function_call = response.choices[0].message.function_call
|
|
|
|
if function_call:
|
|
function_name = function_call.name
|
|
function_arguments = function_call.arguments
|
|
|
|
if function_name == 'end_my_response':
|
|
context.append({'role': 'function', 'name': function_name, 'content': ''})
|
|
break
|
|
elif function_name == 'end_chat':
|
|
# TODO: add a config option to control whether or not the agent is allowed to do this.
|
|
print(colored('The agent has terminated the connection.', 'red', attrs=['bold']))
|
|
sys.exit(1)
|
|
|
|
print(colored(f'{function_name}("{json.dumps(json.loads(function_arguments), indent=2)}")' + '\n', 'yellow'))
|
|
|
|
if function_name not in VALID_FUNCS:
|
|
context.append({'role': 'system', 'content': f'"{function_name}" is not a valid function. Valid functions are {VALID_FUNCS}.'})
|
|
print(colored(f'Attempted to use invalid function {function_name}("{function_arguments}")' + '\n', 'yellow'))
|
|
else:
|
|
# TODO: don't hardcode this
|
|
if function_name == 'run_bash':
|
|
command_output = func_run_bash(function_arguments)
|
|
result_to_ai = {
|
|
'function': function_name,
|
|
'input': function_arguments,
|
|
'stdout': command_output[0],
|
|
'stderr': command_output[1],
|
|
'return_code': command_output[2]
|
|
}
|
|
context.append({'role': 'function', 'name': function_name, 'content': json.dumps({'args': function_arguments, 'result': result_to_ai}, separators=(',', ':'))})
|
|
elif function_name == 'search_google':
|
|
command_output = search_google(json.loads(function_arguments)['query'])
|
|
context.append({'role': 'function', 'name': function_name, 'content': jsonify_anything({'args': function_arguments, 'result': command_output})})
|
|
elif function_name == 'search_google_maps':
|
|
args = json.loads(function_arguments)
|
|
del args['reasoning']
|
|
command_output = search_google_maps(**args)
|
|
context.append({'role': 'function', 'name': function_name, 'content': jsonify_anything({'args': function_arguments, 'result': command_output})})
|
|
elif function_name == 'search_google_news':
|
|
command_output = search_google_news(json.loads(function_arguments)['query'])
|
|
context.append({'role': 'function', 'name': function_name, 'content': jsonify_anything({'args': json.loads(function_arguments), 'result': command_output})})
|
|
# Restart the loop to let the agent decide what to do next.
|
|
else:
|
|
response_text = response.choices[0].message.content
|
|
if response_text == context[-1]['content']:
|
|
# Try to skip duplicate messages.
|
|
break
|
|
|
|
# Sometimes the agent will get confused and send "end_my_response" in the message body. We know what he means.
|
|
end_my_response = True if 'end_my_response' in response_text or not response_text.strip() else False
|
|
response_text = re.sub(r'\n*end_my_response', '', response_text)
|
|
|
|
context.append({'role': 'assistant', 'content': response_text})
|
|
|
|
# We need to print each line individually since the colored text doesn't support the "\n" character
|
|
lines = response_text.split('\n')
|
|
for line in lines:
|
|
print(colored(line, 'blue'))
|
|
print()
|
|
|
|
if end_my_response:
|
|
break
|
|
i += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|