app.py

from openai import OpenAI
import gradio as gr
import os
 
client = OpenAI(api_key=os.environ['openAIToken'])
assistantId = os.environ['assistantId']
initial_message = os.environ['initialMessage']

# client = OpenAI(api_key="sk-proj-4GPZmEnOFsZx0fZeQ9B3T3BlbkFJkXHmGCuW0rGmK0RAavtY")
# assistantId = "asst_Rce5MptxfyEwp4NjImSDvXGD"
# initial_message = "Hello. Let's start defining new facility. Also provide options."

def createThread(sessionStorage):
  streaming_thread = client.beta.threads.create()
  sessionStorage.append(streaming_thread.id)
  return sessionStorage

def addMessage(text, threadId):
  print("User message: ", text)
  return client.beta.threads.messages.create(
    thread_id=threadId,
    role="user",
    content=text
  )

list_of_suggestions = []
latest_suggestions = ""

def handle_suggestions(string_of_suggestions):
  local_message = None
  parts = string_of_suggestions.split('#s#')
  list_of_suggestions = [x.strip('"') for x in parts[0].strip('][').split('", ') if x.strip('"')]
  list_of_suggestions = [ x for x in list_of_suggestions if x]
  if len(parts) > 1:
    local_message = parts[1]
  return list_of_suggestions, local_message

def create_suggestions_list(suggestions):
  update_show = [gr.update(visible=True, value=w) for w in suggestions]
  update_hide = [gr.update(visible=False, value="") for _ in range(6-len(suggestions))]
  return update_show + update_hide

btn_list = []
list_of_suggestions = ""

CSS ="""
.contain { display: flex; flex-direction: column; }
.gradio-container { height: 100vh !important; }
#chatbot { flex-grow: 1; overflow: auto;}
"""

def initiate_chatting(storage):
  threadId = storage[0]
  addMessage(initial_message, threadId)
  global list_of_suggestions
  list_of_suggestions = []
  string_of_suggestions = ""
  is_loading_suggestions = False
  is_it_first_response = True
  response_message = ""
  with client.beta.threads.runs.stream(
    thread_id=threadId,
    assistant_id=assistantId,
  ) as stream:
    for text in stream.text_deltas:
      local_message = None
      if "[" in text and is_it_first_response:
        is_loading_suggestions = True

      is_it_first_response = False

      if is_loading_suggestions != True:
        local_message = text
      else:
        string_of_suggestions = string_of_suggestions + text
        if "#s#" in string_of_suggestions:
          is_loading_suggestions = False
          list_of_suggestions, local_message = handle_suggestions(string_of_suggestions)
      if local_message is not None:
        response_message += local_message 
    stream.until_done()
  
  return [[None, response_message]]

def createDemo():
  with gr.Blocks(css=CSS, fill_height=True) as demo:
    storage = gr.State([])
    chatbot = gr.Chatbot(label="Facility managment bot", elem_id="chatbot")
    with gr.Row():
      for i in range(6):
          btn = gr.Button(visible=False)
          btn_list.append(btn)
    msg = gr.Textbox(label="Answer", interactive=False)
    with gr.Row():
      clear_btn = gr.ClearButton(chatbot)
      btn = gr.Button("Submit")


    def user(user_message, history):
        return "", history + [[user_message, None]]

    def respond(chat_history, storage):
      global btn_list
      message = chat_history[-1][0]
      threadId = storage[0]
      print("Responding for threadId: ", threadId)
      chat_history[-1][1] = ""
      addMessage(message, threadId)
      global list_of_suggestions
      list_of_suggestions = []
      string_of_suggestions = ""
      is_loading_suggestions = False
      is_it_first_response = True
      try:
        with client.beta.threads.runs.stream(
          thread_id=threadId,
          assistant_id=assistantId,
        ) as stream:
          for text in stream.text_deltas:
            print(text, end="")
            local_message = None
            if "[" in text and is_it_first_response:
              is_loading_suggestions = True

            is_it_first_response = False

            if is_loading_suggestions != True:
              local_message = text
            else:
              string_of_suggestions = string_of_suggestions + text
              if "#s#" in string_of_suggestions:
                is_loading_suggestions = False
                list_of_suggestions, local_message = handle_suggestions(string_of_suggestions)
            if local_message is not None:
              chat_history[-1][1] += local_message 
              yield {chatbot: chat_history}
          stream.until_done()
        print("")
      except:
        list_of_suggestions = []
        chat_history[-1][1] = "Error occured during processing your message. Please try again"
        yield {chatbot: chat_history}

    def update_suggestions():
      global list_of_suggestions
      btn_list = create_suggestions_list(list_of_suggestions)
      return btn_list
    
    def hide_suggestions(): 
      return [gr.update(visible=False, value="") for _ in range(6)]

    def disable_msg():
      message_box = gr.Textbox(value=None, interactive=False)
      return message_box
    
    def enable_msg():
      message_box = gr.Textbox(value=None, interactive=True)
      return message_box

    add_user_message_flow = [user, [msg,chatbot],  [msg,chatbot]]
    chat_response_flow = [respond, [chatbot, storage], [chatbot]]
    update_suggestions_flow = [update_suggestions, None, btn_list]
    hide_suggestions_flow = [hide_suggestions, None, btn_list]
    disable_msg_flow = [disable_msg, None, msg] 
    enable_msg_flow = [enable_msg, None, msg]

    btn.click(*add_user_message_flow
        ).then(*hide_suggestions_flow
        ).then(*disable_msg_flow
        ).then(*chat_response_flow
        ).then(*update_suggestions_flow
        ).then(*enable_msg_flow)
    msg.submit(*add_user_message_flow
        ).then(*hide_suggestions_flow
        ).then(*disable_msg_flow
        ).then(*chat_response_flow
        ).then(*update_suggestions_flow
        ).then(*enable_msg_flow)
    for sug_btn in btn_list:
      add_suggestion_message_flow = [user, [sug_btn, chatbot], [msg, chatbot]]
      sug_btn.click(*add_suggestion_message_flow
              ).then(*hide_suggestions_flow
              ).then(*disable_msg_flow
              ).then(*chat_response_flow
              ).then(*update_suggestions_flow
              ).then(*enable_msg_flow)
    
    demo.load(createThread, inputs=storage, outputs=storage
        ).then(initiate_chatting, inputs=storage, outputs=chatbot
        ).then(*update_suggestions_flow
        ).then(*enable_msg_flow)
  return demo

if __name__ == "__main__":
  demo = createDemo()
  demo.launch()