from openai import OpenAI
import gradio as gr
import re
import time 


client = OpenAI() 
assistant = client.beta.assistants.retrieve("asst_G71jWtWDGyurCE6g5C7ABzUx")
thread = client.beta.threads.create()

flightFileLink = 'https://winningchatbot.s3.ap-southeast-1.amazonaws.com/SG-AD-S003' 
tripFileLink= 'https://winningchatbot.s3.ap-southeast-1.amazonaws.com/SG-HR-W013'


def chat(message, history):
    
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), message)
    
    response = client.beta.threads.messages.create(
      thread_id=thread.id,
      role="user",
      content=message
    )
    run = client.beta.threads.runs.create_and_poll(
      thread_id=thread.id,
      assistant_id=assistant.id,
    )

    if run.status == 'completed': 
      messages = client.beta.threads.messages.list(
        thread_id=thread.id
      )
      print(messages)

      reply = messages.data[0].content[0].text.value

      print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), reply)
      
      reply = re.sub(r'\【\d+\:\d+\†source\】', '', reply)

      try: 
        fileId = messages.data[0].content[0].text.annotations[0].file_citation.file_id 

        if 'file-akk9CVNRjnmc3vNk3sV203lN' in fileId:
          reply = reply+ f"\n\n【信息源: SG-AD-S003 韦立新加坡-出差订票和住宿管理方案-修订20240129.pdf】\n({flightFileLink})"
        if 'file-sHwZmYpPK67lsMX4uflSZXD1' in fileId:
          reply = reply+ f"\n\n【信息源: SG-AD-S004 SG-HR-W013 出差制度】\n({tripFileLink})"
      except:
        pass 

      return reply
    else:
      print(run.status)
      return 'running'

gr.ChatInterface(
  fn=chat,
  chatbot= gr.Chatbot(),
  textbox= gr.Textbox(placeholder="开始聊天", container=False, scale=7),
  multimodal = False,
  title="韦立 HR 聊天机器人 (GPT)",
  description="目前所搜集的数据可回答关于韦立的出差订票，住宿管理方案，和出差制度方面问题。",
  # theme="soft",
  examples=["副董事长应该订哪一类机票?", "新加坡的负责订票人员是谁", "在新加坡能订哪些酒店?"],
  cache_examples = False, # will cause hallucination if set to True
  autofocus = True,                 
  retry_btn = None, 
  undo_btn = None,
  clear_btn = None,
).launch()