'''
1. 原生的ChatGPT的中文输出没问题，但是llama index的版本中文输出有问题，而且也可以“继续这个回答”。
1. 可以在terminal中同时运行2个不同版本的ChatGPT，一个原生，一个定制的。python3 xxx.py。
'''

import openai
import gradio as gr
import os

## 这里是huggingface的secret里面存放当前可以用的api key。
openai.api_key = os.environ['my_api_key']


messages = [
    {"role": "system", "content": "You are a helpful and kind AI Assistant."},
]

def chatbot(input):

    history = []
    
    try:
        if input:

            ## no streaming version. 
            # messages.append({"role": "user", "content": input})
            # print(input)
            # chat = openai.ChatCompletion.create(
            #     model="gpt-3.5-turbo", messages=messages
            # )
            # reply = chat.choices[0].message.content
            # messages.append({"role": "assistant", "content": reply})

            ## streaming version. typewriter effect, word by word output. 
            messages.append({"role": "user", "content": input})
            for resp in openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages, stream=True, max_tokens=2048):
                
                # sys.stdout.write(str(resp['choices'][0]['delta'].get('content'))) #! 这个方式仅在terminal中是working的。
                # sys.stdout.flush()
                
                #! 以下内容在Gradio中是working的。
                answer = str(resp['choices'][0]['delta'].get('content'))
                if answer != "None":
                    history.append(answer)
                    result = "".join(history).strip() #* working！
                    yield result
   
    except Exception as e:
        print(e)
        pass
    
    return None



## 如果要优化显示效果，参见： https://gradio.app/theming-guide/
try: 
    inputs = gr.inputs.Textbox(lines=5, label="Prompt提示词")
    outputs = gr.outputs.Textbox(label="ChatGPT")
    interface = gr.Interface(fn=chatbot, inputs=inputs, outputs=outputs, title="石博士的极简版ChatGPT", description="请输入你的提示词",theme=gr.themes.Monochrome(), height=300)
    interface.launch()
except Exception as e:
    print(e)