import g4f
import gradio as gr
from g4f.Provider import (
    Ails,
    You,
    Bing,
    Yqcloud,
    Theb,
    Aichat,
    Bard,
    Vercel,
    Forefront,
    Lockchat,
    Liaobots,
    H2o,
    ChatgptLogin,
    DeepAi,
    GetGpt
)
import os

provider_dict = {
    'Ails': Ails,
    'You': You,
    'Bing': Bing,
    'Yqcloud': Yqcloud,
    'Theb': Theb,
    'Aichat': Aichat,
    'Bard': Bard,
    'Vercel': Vercel,
    'Forefront': Forefront,
    'Lockchat': Lockchat,
    'Liaobots': Liaobots,
    'H2o': H2o,
    'ChatgptLogin': ChatgptLogin,
    'DeepAi': DeepAi,
    'GetGpt': GetGpt
}

with gr.Blocks() as demo:
  chatbot = gr.Chatbot([[None, None]], label='AI')
  msg = gr.Textbox(value="", label='')
  clear = gr.Button("Clear")
  system_msg = gr.Textbox(value="你是一名助手，可以解答问题。", label='系统提示')
  with gr.Row():
    model_name = gr.Dropdown(['gpt-3.5-turbo', 'gpt-4'], value='gpt-3.5-turbo', label='模型')
    provider_name  =  gr.Dropdown(provider_dict.keys(), value='GetGpt', label='提供者')

  def user(user_message, history):
      return gr.update(value="", interactive=False), history + [[user_message, None]]

  def bot(history, model_name, provider_name, system_msg):
      history[-1][1] = ''
      if len(system_msg)>3000:
         system_msg = system_msg[:1500] + system_msg[-1500:]
      prompt = """
      请你仔细阅读以下提示，然后针对用户的话进行回答。
      提示：
      ```
      {}
      ```
      用户最新的话：
      ```
      {}
      ```
      请回答：
      """
      bot_msg = g4f.ChatCompletion.create(model=model_name, 
                                          provider=provider_dict[provider_name], 
                                          messages=[{"role": "user", 
                                                     "content": prompt.format(system_msg,
                                                                              history[-1][0])}],
                                          stream=True)
      for c in bot_msg:
          history[-1][1] += c
          yield history
  
  response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
      bot, [chatbot, model_name, provider_name, system_msg], chatbot
  )
  response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
  clear.click(lambda: None, None, chatbot, queue=False)

demo.title = "AI Chat"
demo.queue()
demo.launch()