import openai
import sys
import os

#api_key = os.getenv('OPENAI_API_KEY')

global_msg = [
    {"role": "system", "content": "你是一位有用的助手"}
]

def get_reply(input_msg):
    msg = {"role": "user", "content": input_msg}
    global_msg.append(msg)
    response = openai.ChatCompletion.create(
        model = "gpt-3.5-turbo",
        top_p = 1,
        frequency_penalty = 0,
        presence_penalty = 0,
        temperature = 1,
        max_tokens = 3000,
        messages = global_msg
    )
    #print(response)
    reply = ''
    choices = response['choices']
    for r in choices:
        reply += r['message']['content']
    return reply


def print_banner():
    banner = '''
       ____ ____ _____ 
      / ___|  _ \_   _|
     | |  _| |_) || |  
     | |_| |  __/ | |  
      \____|_|    |_|  
                   

    '''
    print(banner)
                                            
                                            

def start_robot():
    print_banner()
    print('>>', end='', flush=True)
    while True:
        try:
            msg = sys.stdin.readline().strip()
        except KeyboardInterrupt:
            print('\n再见!')
            break
        if not msg:
            print('>>', end='', flush=True)
            continue
        reply = get_reply(msg)
        reply += '\n>>'
        print('>>' + reply, end='', flush=True)

if __name__ == '__main__':
    start_robot()